82 Commits

Author SHA1 Message Date
Dark-Alex-17 e687d78931 build: Removed unnecessary Language import for Windows systems 2026-05-22 19:04:46 -06:00
Dark-Alex-17 0c2e4df647 feat: LLM node failures propgate up 2026-05-22 18:27:03 -06:00
Dark-Alex-17 6221875f64 build: upgraded to rust v1.95.0 2026-05-22 18:11:01 -06:00
Dark-Alex-17 895b9c27db chore: removed the deprecated haiku 3.5 Claude model 2026-05-22 17:53:49 -06:00
Dark-Alex-17 e661ca2eda docs: Added sharing configurations links in the main README 2026-05-22 17:47:58 -06:00
Dark-Alex-17 7066edd904 feat: Added .install remote tab completions to the REPL 2026-05-22 17:44:16 -06:00
Dark-Alex-17 61bdf29bea feat: feature complete install remote with category selection 2026-05-22 17:00:11 -06:00
Dark-Alex-17 ef39c7d9ff feat: Support to interactively add secrets to Loki that are missing from MCP configs when merging 2026-05-22 16:47:25 -06:00
Dark-Alex-17 e9e46158e7 feat: Added MCP config merging support for remote asset installations 2026-05-22 16:30:45 -06:00
Dark-Alex-17 34dc4b0dce fix: Generified the functions usage of script detection for an executable bit on unix systems 2026-05-22 16:01:28 -06:00
Dark-Alex-17 cd226577e7 feat: install remote now writes files to disk 2026-05-22 15:55:37 -06:00
Dark-Alex-17 b5fc633454 feat: Created basic install_remote functions 2026-05-22 15:33:37 -06:00
Dark-Alex-17 484b18ef16 feat: Created a more comprehensive and immediately useful default config for first runs 2026-05-22 14:16:03 -06:00
Dark-Alex-17 7333046cfe fix: merge required claude code system prompt into instructions 2026-05-22 13:51:45 -06:00
Dark-Alex-17 815f0e5c39 feat: Created an example graph-based agent called deep-research 2026-05-22 12:57:56 -06:00
Dark-Alex-17 dacccbfcf7 feat: Improved coder agent that is now a graph-based agent 2026-05-22 12:57:12 -06:00
Dark-Alex-17 5370637274 docs: Removed slightly-confusing wording in the README 2026-05-22 12:56:49 -06:00
Dark-Alex-17 e6da252a5a feat: Removed indicatif spinners. The UX just won't stop clobbering for parallel graph nodes 2026-05-22 12:56:04 -06:00
Dark-Alex-17 4aaff21f45 fix: updated argc argument passing in run-tool and run-agent scripts 2026-05-21 17:06:20 -06:00
Dark-Alex-17 2678afe02b docs: updated the graph.example.yaml to document the agent environment variables. 2026-05-21 13:29:38 -06:00
Dark-Alex-17 558b764db8 feat: Added agent variables support for graph agents and improved script executor to use the same environment variables as normal agent tool calling for further flexibility 2026-05-21 13:27:33 -06:00
Dark-Alex-17 0bb312a85c feat: Improved UX with colored spinners for parallel graph agents and no clobbering outputs for sub-agents 2026-05-21 13:00:44 -06:00
Dark-Alex-17 d81d233527 feat: created new graph-based deep-research agent 2026-05-21 11:27:55 -06:00
Dark-Alex-17 597f823bdf fmt: cleaned up graph implementation 2026-05-21 11:27:29 -06:00
Dark-Alex-17 81c037515e feat: improved UX for parallel graph execution 2026-05-20 18:54:20 -06:00
Dark-Alex-17 3c7d19da07 fix: Added additional graph validation for parallel reads and writes with dependencies between nodes states 2026-05-20 17:35:33 -06:00
Dark-Alex-17 4536d00067 docs: created an example graph agent configuration 2026-05-20 16:54:34 -06:00
Dark-Alex-17 98d16d9a56 fix: bug in next_single method and improved outcome handling for LLM node execution 2026-05-20 16:27:25 -06:00
Dark-Alex-17 26de81e84e test: implemented integration tests for the parallel frontier-based graph scheduling 2026-05-20 16:09:07 -06:00
Dark-Alex-17 20c28b55d5 feat: added branch progress tracker for better visualization of parallel graph super-steps 2026-05-20 15:50:38 -06:00
Dark-Alex-17 7d6f1dda26 feat: Removed the jira-helper agent and replaced it with the atlassian role 2026-05-20 15:38:51 -06:00
Dark-Alex-17 9a061944ae feat: created the RenderMode enum to suppress stdout streaming during parallel graph super-steps 2026-05-20 15:32:03 -06:00
Dark-Alex-17 1f50af0974 feat: Full support for map node types 2026-05-20 15:15:58 -06:00
Dark-Alex-17 bdacf9fc78 feat: implemented the frontier-based scheduling for the graph executor with simplified state management (gotta love .clone) 2026-05-20 13:48:55 -06:00
Dark-Alex-17 a9f2a5edc2 feat: validation support for parallel graph execution; restricted map nodes to only run for nodes without next targets and not supporting chained map nodes 2026-05-20 12:50:29 -06:00
Dark-Alex-17 2df8b1a541 fix: inline RAG bug when globbing files by extension without subdirectory globbing 2026-05-20 12:22:21 -06:00
Dark-Alex-17 de055bf8a4 feat: created the staging area for state merges per super-step and created the built-in reducers (and their application) for the state merge phase of a super step 2026-05-20 12:16:14 -06:00
Dark-Alex-17 8fb0eece4b feat: scaffolding work for fan-out nodes for parallel branch execution support and stubbed out Map node types 2026-05-20 11:37:23 -06:00
Dark-Alex-17 ba03c3037d style: applied formatting to the new update feature 2026-05-19 14:44:15 -06:00
Dark-Alex-17 afa0e4af67 feat: Loki can now update itself via .update and --update commands 2026-05-19 14:29:44 -06:00
Dark-Alex-17 5a9a00bc6f build: updated dependencies to the latest versions and removed unused dependencies 2026-05-19 13:03:31 -06:00
Dark-Alex-17 e7bb668ac7 fix: update the estimate_token_length function to use the standard word count method 2026-05-19 12:25:53 -06:00
Dark-Alex-17 04498b96ec fix: removed unnecessary regenerate logic for sessions and use the same logic for all contexts; prevents a panic on empty message list 2026-05-19 11:46:37 -06:00
Dark-Alex-17 eb2843d38a build: upgraded to the most recent version of reqwest 2026-05-19 11:05:40 -06:00
Dark-Alex-17 696ce03ee4 feat: added a .edit command for editing the MCP configuration file 2026-05-18 15:14:22 -06:00
Dark-Alex-17 a3d67bfbf7 feat: Created a new .install command to install bundled assets on-demand 2026-05-18 14:59:02 -06:00
Dark-Alex-17 5bd0766a60 style: Cleaned up all graph agent code 2026-05-18 13:46:52 -06:00
Dark-Alex-17 35e1b14843 fix: error when users try to start a session on a graph agent 2026-05-18 12:55:17 -06:00
Dark-Alex-17 503c9b4699 feat: migrated llm node validation to graph loading time instead of graph runtime 2026-05-18 11:51:47 -06:00
Dark-Alex-17 7a8b09542d feat: ripped out user input timeout scaffolding for approval and input node types; implementation can't be done cleanly 2026-05-18 11:32:34 -06:00
Dark-Alex-17 da5cd21c1c test: added additional test coverage to graph components 2026-05-18 10:08:36 -06:00
Dark-Alex-17 27fcb1fc15 docs: Updated README and created graph.example.yaml spec 2026-05-15 17:37:54 -06:00
Dark-Alex-17 e292c414c5 feat: added additional support for all RAG-configuration fields in RAG nodes 2026-05-15 16:38:52 -06:00
Dark-Alex-17 8a2f18204f feat: initial support for RAG nodes in the graph execution system 2026-05-15 14:11:23 -06:00
Dark-Alex-17 c70ac98223 feat: implemented structured logging for graph execution 2026-05-15 13:17:42 -06:00
Dark-Alex-17 249d1fc881 feat: merged normal agent config and graph agent configs into one file (either/or) 2026-05-15 12:57:08 -06:00
Dark-Alex-17 3f4fd91b3f fix: added on_other field for approval nodes so users can specify an alternative free-text target when none of the options match what they want 2026-05-14 16:35:08 -06:00
Dark-Alex-17 48c52b5829 feat: added structured-output extraction for llm and agent nodes 2026-05-14 15:36:10 -06:00
Dark-Alex-17 f58f751c59 fix: accidentally added back in full agent tools on LLM nodes 2026-05-14 14:39:08 -06:00
Dark-Alex-17 fc7fdc98b4 feat: created full llm node runtime implementation 2026-05-14 14:00:24 -06:00
Dark-Alex-17 f4d7d0fb73 refactor: migrated llm nodes to use Roles to simplify instructions handling and to function like inline roles 2026-05-14 13:24:34 -06:00
Dark-Alex-17 4b38f53488 refactor: migrated the next_node and apply_state_updates logic for LLM nodes into the LlmExecutor 2026-05-14 12:08:55 -06:00
Dark-Alex-17 186422ff58 feat: scaffolded together the initial llm node type and its executor 2026-05-14 11:57:18 -06:00
Dark-Alex-17 9bc4f8b621 feat: wired together graph execution and agent graph dispatch 2026-05-14 11:10:45 -06:00
Dark-Alex-17 84497d3d65 feat: implemented support for the graph executor 2026-05-13 14:29:45 -06:00
Dark-Alex-17 3ea9116a23 feat: created the approval node executor and the input node executor for user interaction 2026-05-13 14:08:44 -06:00
Dark-Alex-17 bfcd73c32a feat: Added initial support for native Loki agent nodes in the graph-based agent system 2026-05-13 13:21:45 -06:00
Dark-Alex-17 3cd3ba55ff feat: Added direct script invocation support for graph-based agents 2026-05-13 12:35:10 -06:00
Dark-Alex-17 3535edba79 feat: Added graph validation 2026-05-13 10:18:51 -06:00
Dark-Alex-17 bf0343e245 feat: Implemented state management for agent graphs 2026-05-13 09:18:38 -06:00
Dark-Alex-17 b001ae4c18 feat: initial agent graph scaffolding 2026-05-12 14:13:03 -06:00
Dark-Alex-17 9ce088a530 fix: Improve the coder agent's usage of tools 2026-05-11 15:03:15 -06:00
Dark-Alex-17 16f3f71188 fix: make the agent__collect escalation-aware so it doesn't freeze on sub-agent escalations 2026-05-11 13:57:02 -06:00
Dark-Alex-17 0af5fa02f9 fmt: Applied uniform formatting across all files 2026-05-08 15:52:12 -06:00
Dark-Alex-17 d6a0676264 docs: Updated example configurations to link to the new Wiki-based documentation 2026-05-08 15:51:11 -06:00
Dark-Alex-17 b582bab17c fix: check for an existing session before starting up MCP servers when switching to a role 2026-05-08 12:28:24 -06:00
Dark-Alex-17 a8732c63d6 fix: do not switch to agent if a session is active. 2026-05-08 12:15:01 -06:00
Dark-Alex-17 389d0b768f fix: Do not append todo instructions when function calling is disabled 2026-05-08 12:06:07 -06:00
Dark-Alex-17 70a251a7e2 feat: add auto-continue support to all contexts 2026-05-08 12:02:10 -06:00
Dark-Alex-17 462f136596 feat: dynamic tab completions now show the sessions for a given agent instead of only listing global sessions 2026-05-07 15:23:50 -06:00
Dark-Alex-17 bf9d7d750e fix: a bug in the dynamic completions because the crate name is loki-ai but the binary is named loki 2026-05-07 14:08:54 -06:00
Alex Clarke 540ec648c9 Merge pull request #11 from Dark-Alex-17/config-refactor
Decompose God-Config struct into focused state architecture with MCP SSE support and comprehensive tests
2026-05-07 13:50:49 -06:00
85 changed files with 15368 additions and 1131 deletions
Generated
+776 -357
View File
File diff suppressed because it is too large Load Diff
+15 -17
View File
@@ -10,7 +10,7 @@ repository = "https://github.com/Dark-Alex-17/loki"
categories = ["command-line-utilities"]
readme = "README.md"
license = "MIT"
rust-version = "1.89.0"
rust-version = "1.95.0"
exclude = [".github", "CONTRIBUTING.md"]
[dependencies]
@@ -22,7 +22,7 @@ dunce = "1.0.5"
futures-util = "0.3.29"
inquire = "0.9.4"
is-terminal = "0.4.9"
reedline = "0.46.0"
reedline = "0.47.0"
serde = { version = "1.0.152", features = ["derive"] }
serde_json = { version = "1.0.93", features = ["preserve_order"] }
serde_yaml = "0.9.17"
@@ -34,10 +34,6 @@ tokio = { version = "1.34.0", features = [
"rt-multi-thread",
"full",
] }
tokio-graceful = "0.2.2"
tokio-stream = { version = "0.1.15", default-features = false, features = [
"sync",
] }
crossterm = "0.29.0"
chrono = "0.4.23"
bincode = { version = "2.0.0", features = [
@@ -51,7 +47,7 @@ nu-ansi-term = "0.50.0"
async-trait = "0.1.74"
textwrap = "0.16.0"
ansi_colours = "1.2.2"
reqwest-eventsource = "0.6.0"
eventsource-stream = "0.2.3"
log = "0.4.28"
log4rs = { version = "1.4.0", features = ["file_appender"] }
shell-words = "1.1.0"
@@ -59,20 +55,14 @@ sha2 = "0.10.8"
unicode-width = "0.2.0"
async-recursion = "1.1.1"
http = "1.1.0"
http-body-util = "0.1"
hyper = { version = "1.0", features = ["full"] }
hyper-util = { version = "0.1", features = ["server-auto", "client-legacy"] }
time = { version = "0.3.36", features = ["macros"] }
indexmap = { version = "2.2.6", features = ["serde"] }
hmac = "0.12.1"
aws-smithy-eventstream = "0.60.4"
urlencoding = "2.1.3"
unicode-segmentation = "1.11.0"
json-patch = { version = "4.0.0", default-features = false }
bitflags = "2.5.0"
path-absolutize = "3.1.1"
hnsw_rs = "0.3.0"
rayon = "1.10.0"
uuid = { version = "1.9.1", features = ["v4"] }
scraper = { version = "0.23.1", default-features = false, features = [
"deterministic",
@@ -97,7 +87,6 @@ rmcp = { version = "1.5.0", features = [
] }
num_cpus = "1.17.0"
tree-sitter = "0.26.8"
tree-sitter-language = "0.1"
tree-sitter-python = "0.25.0"
tree-sitter-typescript = "0.23"
colored = "3.0.0"
@@ -107,15 +96,24 @@ clap_complete_nushell = "4.5.9"
open = "5"
rand = { version = "0.10.0", features = ["default"] }
url = "2.5.8"
self_update = { version = "0.44", default-features = false, features = [
"reqwest",
"rustls",
"archive-tar",
"compression-flate2",
"archive-zip",
"compression-zip-deflate",
] }
[dependencies.reqwest]
version = "0.12.0"
version = "0.13.3"
features = [
"json",
"multipart",
"stream",
"form",
"socks",
"rustls-tls",
"rustls-tls-native-roots",
"rustls",
]
default-features = false
+30 -15
View File
@@ -10,7 +10,8 @@ Loki is an all-in-one, batteries-included, LLM CLI tool featuring Shell Assistan
Agents, and More.
It is designed to include a number of useful agents, roles, macros, and more so users can get up and running with Loki
in as little time as possible.
in as little time as possible. You can also install entire bundles of agents, roles, macros, tools, and MCP servers from
any git repository — see [Sharing Configurations](#sharing-configurations).
![Agent example](https://raw.githubusercontent.com/wiki/Dark-Alex-17/loki/images/agents/sql.gif)
@@ -20,6 +21,7 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
* [AIChat Migration Guide](https://github.com/Dark-Alex-17/loki/wiki/AIChat-Migration): Coming from AIChat? Follow the migration guide to get started.
* [Installation](#install): Install Loki
* [Getting Started](#getting-started): Get started with Loki by doing first-run setup steps.
* [Sharing Configurations](https://github.com/Dark-Alex-17/loki/wiki/Sharing-Configurations): Install bundles of agents, roles, macros, tools, and MCP servers from any git repo, and share your own.
* [REPL](https://github.com/Dark-Alex-17/loki/wiki/REPL): Interactive Read-Eval-Print Loop for conversational interactions with LLMs and Loki.
* [Custom REPL Prompt](https://github.com/Dark-Alex-17/loki/wiki/REPL-Prompt): Customize the REPL prompt to provide useful contextual information.
* [Vault](https://github.com/Dark-Alex-17/loki/wiki/Vault): Securely store and manage sensitive information such as API keys and credentials.
@@ -36,7 +38,8 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
* [Sessions](https://github.com/Dark-Alex-17/loki/wiki/Sessions): Manage and persist conversational contexts and settings across multiple interactions.
* [Roles](https://github.com/Dark-Alex-17/loki/wiki/Roles): Customize model behavior for specific tasks or domains.
* [Agents](https://github.com/Dark-Alex-17/loki/wiki/Agents): Leverage AI agents to perform complex tasks and workflows, including sub-agent spawning, teammate messaging, and user interaction tools.
* [Todo System](https://github.com/Dark-Alex-17/loki/wiki/TODO-System): Built-in task tracking for improved agent reliability with smaller models.
* [Graph Agents](https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents): Define an agent as a declarative, YAML-driven workflow. A directed graph of typed nodes (LLM calls, scripts, approvals, user input, RAG retrieval, sub-agent spawns).
* [Todo System](https://github.com/Dark-Alex-17/loki/wiki/TODO-System): Built-in task tracking for improved LLM reliability with smaller models.
* [Environment Variables](https://github.com/Dark-Alex-17/loki/wiki/Environment-Variables): Override and customize your Loki configuration at runtime with environment variables.
* [Client Configurations](https://github.com/Dark-Alex-17/loki/wiki/Clients): Configuration instructions for various LLM providers.
* [Authentication (API Key & OAuth)](https://github.com/Dark-Alex-17/loki/wiki/Clients#authentication): Authenticate with API keys or OAuth for subscription-based access.
@@ -48,16 +51,6 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
Loki requires the following tools to be installed on your system:
* [jq](https://github.com/jqlang/jq)
* `brew install jq`
* [jira (optional)](https://github.com/ankitpokhrel/jira-cli/wiki/Installation) (For the `query_jira_issues` tool)
* `brew tap ankitpokhrel/jira-cli && brew install jira-cli`
* You'll need to [create a JIRA API token](https://id.atlassian.com/manage-profile/security/api-tokens) for authentication
* Then, save it as an environment variable to your shell profile:
```sh
# ~/.bashrc or ~/.zshrc
export JIRA_API_TOKEN="your_jira_api_token_here"
```
* Then run `jira init`, select installation type as `cloud`, and provide the required details to generate a config
file for the Jira CLI.
* [usql](https://github.com/xo/usql) (For the `sql` agent)
* `brew install xo/xo/usql`
* [docker](https://docs.docker.com/engine/install/)
@@ -65,7 +58,7 @@ Loki requires the following tools to be installed on your system:
* `curl -LsSf https://astral.sh/uv/install.sh | sh`
These tools are used to provide various functionalities within Loki, such as document processing, JSON manipulation,
interaction with Jira, and they are used within agents and tools.
etc., and they are used within agents and tools.
## Install
@@ -137,6 +130,29 @@ To use a binary from the releases page on Linux/MacOS, do the following:
3. Extract the binary with `tar -C /usr/local/bin -xzf loki-<arch>.tar.gz` (Note: This may require `sudo`)
4. Now you can run `loki`!
## Updating
Loki can update itself in place to the latest GitHub release. Run `loki --update`
for the newest release, or `loki --update v0.4.0` for a specific version:
```shell
loki --update
loki --update v0.4.0
```
The same is available from within the REPL via `.update` and `.update v0.4.0`.
If Loki was installed with a package manager, prefer that package manager so its
records stay in sync with the binary on disk; i.e. `brew upgrade loki` for Homebrew,
or `cargo install --locked loki-ai` for Cargo.
When Loki detects a package-manager install it prints a warning and asks for
confirmation. In a non-interactive shell (no TTY), pass `--force` to update
anyway:
```shell
loki --update --force
```
## Getting Started
After installation, you can generate the configuration files and directories by simply running:
@@ -160,12 +176,11 @@ subscribers, Google Gemini), you can authenticate with your existing subscriptio
# In your config.yaml
clients:
- type: claude
name: my-claude-oauth
auth: oauth # Indicate you want to authenticate with OAuth instead of an API key
```
```sh
loki --authenticate my-claude-oauth
loki --authenticate claude
# Or via the REPL: .authenticate
```
+63 -21
View File
@@ -1,40 +1,82 @@
# Coder
An AI agent that assists you with your coding tasks.
A graph-based implementation agent. Plans, implements, and runs build +
tests in a bounded fix-loop until verified. Designed to be delegated to by
the **[Sisyphus](../sisyphus/README.md)** agent.
This agent is designed to be delegated to by the **[Sisyphus](../sisyphus/README.md)** agent to implement code specifications. Sisyphus
acts as the coordinator/architect, while Coder handles the implementation details.
Coder is a [graph agent](https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents): its workflow is
defined declaratively in `graph.yaml`, with verification and the
implement-fix loop enforced as graph edges rather than prose.
## Features
## Workflow
- 🏗️ Intelligent project structure creation and management
- 🖼️ Convert screenshots into clean, functional code
- 📁 Comprehensive file system operations (create folders, files, read/write files)
- 🧐 Advanced code analysis and improvement suggestions
- 📊 Precise diff-based file editing for controlled code modifications
```
analyze_request (llm + output_schema) plan + complexity extraction
route_complexity (script) opt-out approval gate (complexity ≥ 7)
gate_approval (approval, optional)
implement (llm + fs tools) actual file edits
verify_build (script)
verify_tests (script)
fix_loop_gate (script) back-edge to implement (bounded)
end_success / end_rejected / end_failure
```
It can also be used as a standalone tool for direct coding assistance.
End nodes emit one of three sentinel outcomes for the caller:
## Pro-Tip: Use an IDE MCP Server for Improved Performance
Many modern IDEs now include MCP servers that let LLMs perform operations within the IDE itself and use IDE tools. Using
an IDE's MCP server dramatically improves the performance of coding agents. So if you have an IDE, try adding that MCP
server to your config (see the [MCP Server docs](../../../docs/function-calling/MCP-SERVERS.md) to see how to configure
them), and modify the agent definition to look like this:
- `CODER_COMPLETE` — build and tests passed.
- `CODER_REJECTED` — user rejected the plan at the approval gate.
- `CODER_FAILED` — fix-loop exhausted; build/tests still failing.
## Tuning
The agent's `project_dir` is exposed via the standard `variables:` block,
so it accepts the runtime override flag:
```sh
# Invoke from inside the project (project_dir defaults to ".")
cd /path/to/your/project
loki -a coder "Add a foo() function..."
# Or invoke from anywhere with an explicit override
loki -a coder --agent-variable project_dir /path/to/your/project "Add..."
```
`graph.yaml` `initial_state` exposes:
- `max_fix_attempts` (default `3`) — fix-loop budget before `end_failure`.
Environment overrides honored by the script nodes:
- `BUILD_CMD` — skip project-type detection for the build/check command.
- `TEST_CMD` — skip detection for tests.
- `CODER_AUTOAPPROVE=1` — bypass the approval gate (for non-interactive runs
where complexity might trip the gate).
## Pro-Tip: IDE MCP Server
Modern IDEs (JetBrains, VS Code, Cursor, Zed, etc.) expose MCP servers
that let LLMs use IDE tools directly. To wire one in, edit `graph.yaml`:
```yaml
# ...
mcp_servers:
- jetbrains # The name of your configured IDE MCP server
- your-ide-mcp-server
global_tools:
# Keep useful read-only tools for reading files in other non-project directories
# Keep read-only fs tools for files outside the IDE project
- fs_read.sh
- fs_grep.sh
- fs_glob.sh
# - fs_write.sh
# - fs_patch.sh
- execute_command.sh
# ...
```
Then add the MCP server's write/patch tools to the `implement` node's
`tools:` whitelist.
-129
View File
@@ -1,129 +0,0 @@
name: coder
description: Implementation agent - writes code, follows patterns, verifies with builds
version: 1.0.0
temperature: 0.1
auto_continue: true
max_auto_continues: 15
inject_todo_instructions: true
variables:
- name: project_dir
description: Project directory to work in
default: '.'
- name: auto_confirm
description: Auto-confirm command execution
default: '1'
global_tools:
- fs_read.sh
- fs_grep.sh
- fs_glob.sh
- fs_write.sh
- fs_patch.sh
- execute_command.sh
instructions: |
You are a senior engineer. You write code that works on the first try.
## Your Mission
Given an implementation task:
1. Check for orchestrator context first (see below)
2. Fill gaps only. Read files NOT already covered in context
3. Write the code (using tools, NOT chat output)
4. Verify it compiles/builds
5. Signal completion with a summary
## Using Orchestrator Context (IMPORTANT)
When spawned by sisyphus, your prompt will often contain a `<context>` block
with prior findings: file paths, code patterns, and conventions discovered by
explore agents.
**If context is provided:**
1. Use it as your primary reference. Don't re-read files already summarized
2. Follow the code patterns shown. Snippets in context ARE the style guide
3. Read the referenced files ONLY IF you need more detail (e.g. full function
signature, import list, or adjacent code not included in the snippet)
4. If context includes a "Conventions" section, follow it exactly
**If context is NOT provided or is too vague to act on:**
Fall back to self-exploration: grep for similar files, read 1-2 examples,
match their style.
**Never ignore provided context.** It represents work already done upstream.
## Todo System
For multi-file changes:
1. `todo__init` with the implementation goal
2. `todo__add` for each file to create/modify
3. Implement each, calling `todo__done` immediately after
## Writing Code
**CRITICAL**: Write code using `write_file` tool, NEVER paste code in chat.
Correct:
```
write_file --path "src/user.rs" --content "pub struct User { ... }"
```
Wrong:
```
Here's the implementation:
\`\`\`rust
pub struct User { ... }
\`\`\`
```
## File Reading Strategy (IMPORTANT - minimize token usage)
1. **Use grep to find relevant code** - `fs_grep --pattern "fn handle_request" --include "*.rs"` finds where things are
2. **Read only what you need** - `fs_read --path "src/main.rs" --offset 50 --limit 30` reads lines 50-79
3. **Never cat entire large files** - If 500+ lines, read the relevant section after grepping for it
4. **Use glob to find files** - `fs_glob --pattern "*.rs" --path src/` discovers files by name
## Pattern Matching
Before writing ANY file:
1. Find a similar existing file (use `fs_grep` to locate, then `fs_read` to examine)
2. Match its style: imports, naming, structure
3. Follow the same patterns exactly
## Verification
After writing files:
1. Run `verify_build` to check compilation
2. If it fails, fix the error (minimal change)
3. Don't move on until build passes
## Completion Signal
When done, end your response with a summary so the parent agent knows what happened:
```
CODER_COMPLETE: [summary of what was implemented, which files were created/modified, and build status]
```
Or if something went wrong:
```
CODER_FAILED: [what went wrong]
```
## Rules
1. **Write code via tools** - Never output code to chat
2. **Follow patterns** - Read existing files first
3. **Verify builds** - Don't finish without checking
4. **Minimal fixes** - If build fails, fix precisely
5. **No refactoring** - Only implement what's asked
## Context
- Project: {{project_dir}}
- CWD: {{__cwd__}}
- Shell: {{__shell__}}
## Available tools:
{{__tools__}}
+278
View File
@@ -0,0 +1,278 @@
name: coder
description: |
Implementation agent. Plans, implements, and runs build + tests in a
bounded fix-loop until verified. Designed to be delegated to by sisyphus.
version: "1.0"
temperature: 0.1
global_tools:
- fs_cat.sh
- fs_ls.sh
- fs_write.sh
- fs_patch.sh
- execute_command.sh
variables:
- name: project_dir
description: |
Absolute path to the project directory. Defaults to "." which is the
directory you invoked `loki` from. Override at runtime with
`loki -a coder --agent-variable project_dir /abs/path "..."`.
default: "."
settings:
max_loop_iterations: 20
log_state_snapshots: true
validate_before_run: true
timeout: 1800
initial_state:
project_dir: ""
fix_attempts: 0
max_fix_attempts: 3
fix_instructions: ""
build_output: ""
tests_output: ""
last_node_output: ""
plan_summary: ""
files_to_modify: []
files_to_create: []
risks: []
complexity_score: 0
start: resolve_paths
nodes:
resolve_paths:
id: resolve_paths
type: script
description: Resolve project_dir to an absolute path from the agent variable
script: scripts/resolve_paths.sh
timeout: 5
fallback: end_failure
analyze_request:
id: analyze_request
type: llm
description: Extract a structured plan and complexity score from the orchestrator's prompt
instructions: |
You are a senior engineer's planning assistant. Read the orchestrator's
request and emit a structured plan. You only plan. You never edit files.
Score complexity from 1 to 10:
1-3: trivial - single file, <=20 lines changed, obvious approach
4-6: moderate - 2-5 files, clear approach, some pattern matching
7-10: complex - multi-component, ambiguous tradeoffs, refactoring,
or wide blast radius
Be specific in `files_to_modify` and `files_to_create`. All paths
MUST be absolute. The project root is {{project_dir}}. Prefer paths
like "{{project_dir}}/src/foo.rs" over "src/foo.rs". The implementer
uses these paths directly with fs_write and fs_patch tools, which
resolve relative paths against the loki invocation directory (NOT
the project dir). Empty arrays are fine if no files in that category.
`risks` is a list of short strings. Anything that could derail the
implementation: unknown dependencies, brittle tests, blast radius,
etc. Empty list is fine.
Project directory: {{project_dir}}
prompt: "{{initial_prompt}}"
tools: []
output_schema:
type: object
properties:
plan_summary:
type: string
description: 1-3 sentences summarizing what will be done
files_to_modify:
type: array
items: {type: string}
files_to_create:
type: array
items: {type: string}
complexity_score:
type: integer
minimum: 1
maximum: 10
risks:
type: array
items: {type: string}
required: [plan_summary, files_to_modify, files_to_create, complexity_score, risks]
state_updates:
last_node_output: "{{output}}"
fallback: end_failure
next: route_complexity
route_complexity:
id: route_complexity
type: script
description: Route to approval gate for complex plans; skip otherwise
script: scripts/route_complexity.sh
timeout: 5
fallback: implement
gate_approval:
id: gate_approval
type: approval
description: Optional human checkpoint for high-complexity plans
question: |
## Plan
{{plan_summary}}
## Files to modify
{{files_to_modify}}
## Files to create
{{files_to_create}}
## Risks
{{risks}}
Complexity: {{complexity_score}}/10
Approve this plan?
options:
- "yes"
- "no"
routes:
"yes": implement
"no": end_rejected
on_other: end_rejected
implement:
id: implement
type: llm
description: Write code via fs tools. Bounded tool-call loop.
instructions: |
You are a senior engineer. Implement the plan by writing code via
tools. Follow existing patterns in the codebase.
## Writing code
1. Use `fs_patch` for surgical edits to existing files.
2. Use `fs_write` for new files or full rewrites.
3. NEVER output code to chat. Always use tools.
4. ALWAYS pass ABSOLUTE paths to fs_write and fs_patch. Relative
paths resolve against the loki invocation directory (not the
project dir), which is rarely what you want. The project root
is {{project_dir}}.
## File reading
1. Use `execute_command` to grep/find:
`execute_command --command "grep -rn 'fn handle_request' --include='*.rs' ."`
`execute_command --command "find . -name '*.rs' -not -path '*/target/*'"`
2. Read only what you need:
`fs_cat --path "src/main.rs" --offset 50 --limit 30`
3. Never read entire large files. Use offset/limit.
4. Use `fs_ls` to list directory contents.
## Pattern matching
Before writing ANY file:
1. Find a similar existing file (grep, then read).
2. Match its style: imports, naming, structure, error handling.
3. Follow the same patterns exactly. Do not invent new ones.
## Fix loop
If the "Fix loop status" section in your user prompt is non-empty,
the previous attempt failed verification. Read the error, identify
the minimal fix, apply it. Do not refactor while fixing.
## Rules
1. Match existing patterns - read examples first.
2. Minimal changes - implement only what's asked.
3. Never suppress errors (`as any`, `@ts-ignore`, `#[allow(...)]`
on unfamiliar lints, etc.).
4. No dead code, no commented-out blocks, no premature abstractions.
5. End your turn when editing is done. The graph runs verification next.
Project directory: {{project_dir}}
prompt: |
## Plan summary
{{plan_summary}}
## Files involved
- Modify: {{files_to_modify}}
- Create: {{files_to_create}}
## Original request from the orchestrator
{{initial_prompt}}
## Fix loop status
{{fix_instructions}}
tools:
- fs_cat
- fs_ls
- fs_write
- fs_patch
- execute_command
max_iterations: 30
state_updates:
last_node_output: "{{output}}"
fallback: end_failure
next: verify_build
verify_build:
id: verify_build
type: script
description: Run the project's check/build command. Routes to verify_tests on success, fix_loop_gate on failure.
script: scripts/verify_build.sh
timeout: 300
fallback: fix_loop_gate
verify_tests:
id: verify_tests
type: script
description: Run the project's test command. Routes to end_success on pass, fix_loop_gate on failure.
script: scripts/verify_tests.sh
timeout: 600
fallback: fix_loop_gate
fix_loop_gate:
id: fix_loop_gate
type: script
description: Budget gate. Loops back to implement with fix_instructions populated, or terminates as end_failure.
script: scripts/fix_loop_gate.sh
timeout: 5
fallback: end_failure
end_success:
id: end_success
type: end
output: |
CODER_COMPLETE
Plan: {{plan_summary}}
Files modified: {{files_to_modify}}
Files created: {{files_to_create}}
Build: passed
Tests: passed
end_rejected:
id: end_rejected
type: end
output: |
CODER_REJECTED
Plan was rejected at the approval gate.
Plan: {{plan_summary}}
end_failure:
id: end_failure
type: end
output: |
CODER_FAILED
Plan: {{plan_summary}}
Attempts: {{fix_attempts}}/{{max_fix_attempts}}
Last node output:
{{last_node_output}}
Last build output:
{{build_output}}
Last tests output:
{{tests_output}}
@@ -0,0 +1,49 @@
#!/usr/bin/env bash
set -euo pipefail
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
state=$(cat "$GRAPH_STATE_FILE")
elif [[ -n "${GRAPH_STATE:-}" ]]; then
state="$GRAPH_STATE"
else
state='{}'
fi
fix_attempts=$(echo "$state" | jq -r '.fix_attempts // 0')
max_fix_attempts=$(echo "$state" | jq -r '.max_fix_attempts // 3')
build_ok=$(echo "$state" | jq -r '.build_ok | if . == null then "true" else (. | tostring) end')
tests_ok=$(echo "$state" | jq -r '.tests_ok | if . == null then "true" else (. | tostring) end')
build_output=$(echo "$state" | jq -r '.build_output // ""')
tests_output=$(echo "$state" | jq -r '.tests_output // ""')
if (( fix_attempts >= max_fix_attempts )); then
jq -nc \
--argjson n "$fix_attempts" \
'{
"fix_attempts": $n,
"_next": "end_failure"
}'
exit 0
fi
next_attempts=$((fix_attempts + 1))
if [[ "$build_ok" != "true" ]]; then
fix_instructions=$(printf '## Fix loop status (attempt %d of %d)\n\nThe previous attempt failed the build.\n\nBuild output:\n```\n%s\n```\n\nIdentify the minimal fix and apply it. Do not refactor.' \
"$next_attempts" "$max_fix_attempts" "$build_output")
elif [[ "$tests_ok" != "true" ]]; then
fix_instructions=$(printf '## Fix loop status (attempt %d of %d)\n\nBuild passed but tests failed.\n\nTest output:\n```\n%s\n```\n\nIdentify the minimal fix and apply it. Do not refactor.' \
"$next_attempts" "$max_fix_attempts" "$tests_output")
else
fix_instructions=$(printf '## Fix loop status (attempt %d of %d)\n\nfix_loop_gate was reached but no failure was detected in state. Re-run the verification step.' \
"$next_attempts" "$max_fix_attempts")
fi
jq -nc \
--argjson n "$next_attempts" \
--arg fi "$fix_instructions" \
'{
"fix_attempts": $n,
"fix_instructions": $fi,
"_next": "implement"
}'
@@ -0,0 +1,12 @@
#!/usr/bin/env bash
set -euo pipefail
project_dir="${LLM_AGENT_VAR_PROJECT_DIR:-.}"
resolved=$(cd "$project_dir" 2>/dev/null && pwd) || resolved="$project_dir"
jq -nc \
--arg pd "$resolved" \
'{
"project_dir": $pd,
"_next": "analyze_request"
}'
@@ -0,0 +1,23 @@
#!/usr/bin/env bash
set -euo pipefail
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
state=$(cat "$GRAPH_STATE_FILE")
elif [[ -n "${GRAPH_STATE:-}" ]]; then
state="$GRAPH_STATE"
else
state='{}'
fi
complexity=$(echo "$state" | jq -r '.complexity_score // 0')
if [[ "${CODER_AUTOAPPROVE:-0}" == "1" ]]; then
jq -nc '{"_next": "implement"}'
exit 0
fi
if (( complexity >= 7 )); then
jq -nc '{"_next": "gate_approval"}'
else
jq -nc '{"_next": "implement"}'
fi
@@ -0,0 +1,55 @@
#!/usr/bin/env bash
set -uo pipefail
# shellcheck disable=SC1091
source "$(dirname "$0")/../../.shared/utils.sh"
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
state=$(cat "$GRAPH_STATE_FILE")
elif [[ -n "${GRAPH_STATE:-}" ]]; then
state="$GRAPH_STATE"
else
state='{}'
fi
project_dir=$(echo "$state" | jq -r '.project_dir // "."')
if [[ -n "${BUILD_CMD:-}" ]]; then
cmd="$BUILD_CMD"
else
project_info=$(detect_project "$project_dir")
cmd=$(echo "$project_info" | jq -r '.check // .build // ""')
fi
if [[ -z "$cmd" || "$cmd" == "null" ]]; then
jq -nc '{
"build_ok": true,
"build_output": "(no build/check command available for this project type)",
"_next": "verify_tests"
}'
exit 0
fi
exit_code=0
output=$(cd "$project_dir" && eval "$cmd" 2>&1) || exit_code=$?
if (( exit_code == 0 )); then
jq -nc \
--arg out "$output" \
--arg cmd "$cmd" \
'{
"build_ok": true,
"build_output": ("Ran: " + $cmd + "\n\n" + $out),
"_next": "verify_tests"
}'
else
jq -nc \
--arg out "$output" \
--arg cmd "$cmd" \
--argjson rc "$exit_code" \
'{
"build_ok": false,
"build_output": ("Ran: " + $cmd + "\nExit code: " + ($rc | tostring) + "\n\n" + $out),
"_next": "fix_loop_gate"
}'
fi
@@ -0,0 +1,55 @@
#!/usr/bin/env bash
set -uo pipefail
# shellcheck disable=SC1091
source "$(dirname "$0")/../../.shared/utils.sh"
if [[ -n "${GRAPH_STATE_FILE:-}" ]]; then
state=$(cat "$GRAPH_STATE_FILE")
elif [[ -n "${GRAPH_STATE:-}" ]]; then
state="$GRAPH_STATE"
else
state='{}'
fi
project_dir=$(echo "$state" | jq -r '.project_dir // "."')
if [[ -n "${TEST_CMD:-}" ]]; then
cmd="$TEST_CMD"
else
project_info=$(detect_project "$project_dir")
cmd=$(echo "$project_info" | jq -r '.test // ""')
fi
if [[ -z "$cmd" || "$cmd" == "null" ]]; then
jq -nc '{
"tests_ok": true,
"tests_output": "(no test command available for this project type)",
"_next": "end_success"
}'
exit 0
fi
exit_code=0
output=$(cd "$project_dir" && eval "$cmd" 2>&1) || exit_code=$?
if (( exit_code == 0 )); then
jq -nc \
--arg out "$output" \
--arg cmd "$cmd" \
'{
"tests_ok": true,
"tests_output": ("Ran: " + $cmd + "\n\n" + $out),
"_next": "end_success"
}'
else
jq -nc \
--arg out "$output" \
--arg cmd "$cmd" \
--argjson rc "$exit_code" \
'{
"tests_ok": false,
"tests_output": ("Ran: " + $cmd + "\nExit code: " + ($rc | tostring) + "\n\n" + $out),
"_next": "fix_loop_gate"
}'
fi
-118
View File
@@ -14,99 +14,6 @@ _project_dir() {
(cd "${dir}" 2>/dev/null && pwd) || echo "${dir}"
}
# Normalize a path to be relative to project root.
# Strips the project_dir prefix if the LLM passes an absolute path.
# Usage: local rel_path; rel_path=$(_normalize_path "/abs/or/rel/path")
_normalize_path() {
local input_path="$1"
local project_dir
project_dir=$(_project_dir)
if [[ "${input_path}" == /* ]]; then
input_path="${input_path#"${project_dir}"/}"
fi
input_path="${input_path#./}"
echo "${input_path}"
}
# @cmd Read a file's contents before modifying
# @option --path! Path to the file (relative to project root)
read_file() {
local file_path
# shellcheck disable=SC2154
file_path=$(_normalize_path "${argc_path}")
local project_dir
project_dir=$(_project_dir)
local full_path="${project_dir}/${file_path}"
if [[ ! -f "${full_path}" ]]; then
warn "File not found: ${file_path}" >> "$LLM_OUTPUT"
return 0
fi
{
info "Reading: ${file_path}"
echo ""
cat "${full_path}"
} >> "$LLM_OUTPUT"
}
# @cmd Write complete file contents
# @option --path! Path for the file (relative to project root)
# @option --content! Complete file contents to write
write_file() {
local file_path
file_path=$(_normalize_path "${argc_path}")
# shellcheck disable=SC2154
local content="${argc_content}"
local project_dir
project_dir=$(_project_dir)
local full_path="${project_dir}/${file_path}"
mkdir -p "$(dirname "${full_path}")"
printf '%s' "${content}" > "${full_path}"
green "Wrote: ${file_path}" >> "$LLM_OUTPUT"
}
# @cmd Find files similar to a given path (for pattern matching)
# @option --path! Path to find similar files for
find_similar_files() {
local file_path
file_path=$(_normalize_path "${argc_path}")
local project_dir
project_dir=$(_project_dir)
local ext="${file_path##*.}"
local dir
dir=$(dirname "${file_path}")
info "Similar files to: ${file_path}" >> "$LLM_OUTPUT"
echo "" >> "$LLM_OUTPUT"
local results
results=$(find "${project_dir}/${dir}" -maxdepth 1 -type f -name "*.${ext}" \
! -name "$(basename "${file_path}")" \
! -name "*test*" \
! -name "*spec*" \
2>/dev/null | sed "s|^${project_dir}/||" | head -3)
if [[ -z "${results}" ]]; then
results=$(find "${project_dir}/src" -type f -name "*.${ext}" \
! -name "*test*" \
! -name "*spec*" \
-not -path '*/target/*' \
2>/dev/null | sed "s|^${project_dir}/||" | head -3)
fi
if [[ -n "${results}" ]]; then
echo "${results}" >> "$LLM_OUTPUT"
else
warn "No similar files found" >> "$LLM_OUTPUT"
fi
}
# @cmd Verify the project builds successfully
verify_build() {
local project_dir
@@ -189,28 +96,3 @@ get_project_structure() {
} >> "$LLM_OUTPUT"
}
# @cmd Search for content in the codebase
# @option --pattern! Pattern to search for
search_code() {
# shellcheck disable=SC2154
local pattern="${argc_pattern}"
local project_dir
project_dir=$(_project_dir)
info "Searching: ${pattern}" >> "$LLM_OUTPUT"
echo "" >> "$LLM_OUTPUT"
local results
results=$(grep -rn "${pattern}" "${project_dir}" 2>/dev/null | \
grep -v '/target/' | \
grep -v '/node_modules/' | \
grep -v '/.git/' | \
sed "s|^${project_dir}/||" | \
head -20) || true
if [[ -n "${results}" ]]; then
echo "${results}" >> "$LLM_OUTPUT"
else
warn "No matches" >> "$LLM_OUTPUT"
fi
}
+274
View File
@@ -0,0 +1,274 @@
# deep-research
A deep web research agent, built as a Loki graph agent. It plans an
investigation, decomposes it into sub-questions researched in
parallel, grounds the work in a local knowledge corpus, vets the
credibility of cited sources, runs a reflexion self-critique loop to
revise weak findings, delegates the final write-up to a focused
sub-agent, checks that the cited sources are reachable, and gates the
result behind human approval.
Unlike a regular agent (which takes a goal and improvises the steps),
this agent runs a fixed graph: every request goes through the same
`plan -> parallel research -> vet -> critique -> synthesize -> verify -> approve`
pipeline.
This agent is also the **canonical reference for the Loki graph
system**: it exercises every node type (`script`, `llm`, `rag`, `map`,
`agent`, `input`, `approval`, `end`) and both static fan-out and
dynamic `map` fan-out. If you are learning how to build a graph
agent, this is the file to read alongside the
[Graph-Agents wiki](https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents).
## Workflow
17 nodes. `->` is the static route; a script node can also route
dynamically via `_next`. The `▶▶` line is a parallel super-step —
those branches run concurrently:
```
parse_request (script) -> bootstrap_research (or -> ask_topic if no topic)
ask_topic (input) -> bootstrap_research
bootstrap_research (script) -> [plan, knowledge_lookup] ▶▶ parallel
plan (llm + output_schema) -> research_each_question
knowledge_lookup (rag) -> research_each_question
research_each_question (map) -> combine_findings (spawns one branch per question)
└─ research_one_question (llm) (atomic; runs N×, joins at map)
combine_findings (script) -> vet_sources
vet_sources (llm + custom tool) -> critique
critique (llm) -> reflexion_gate
reflexion_gate (script) -> synthesize (or -> research_each_question: reflexion loop)
synthesize (agent: report-writer) -> verify_sources
verify_sources (script) -> approve
approve (approval) -> end_accepted ("accept")
-> end_rejected ("reject")
-> incorporate_feedback (any free-form answer)
incorporate_feedback (script) -> research_each_question (the human-feedback loop)
```
### Node-type breakdown
| Type | Nodes |
|---|---|
| `script` (Python) | `parse_request`, `bootstrap_research`, `combine_findings`, `reflexion_gate`, `verify_sources`, `incorporate_feedback` |
| `llm` (tools: `[]`) | `plan`, `critique` |
| `llm` (with tool whitelist) | `research_one_question`, `vet_sources` |
| `rag` | `knowledge_lookup` — local corpus retrieval |
| `map` | `research_each_question` — dynamic fan-out per sub-question |
| `agent` | `synthesize` — spawns the `report-writer` sub-agent |
| `input` | `ask_topic` |
| `approval` | `approve` |
| `end` | `end_accepted`, `end_rejected` |
## Parallel execution
The graph has two parallel super-steps where Loki's BSP scheduler runs
branches concurrently.
**1. Context loading (`plan` ‖ `knowledge_lookup`)** — after
`bootstrap_research`, the LLM planner (which decomposes the topic into
sub-questions) and the RAG retrieval over the local `knowledge/`
corpus run side by side. They write disjoint state keys (`plan` writes
`research_plan` and `questions`; `knowledge_lookup` writes
`local_context` and `local_sources`) so no reducer is needed.
**2. Per-question research (`research_each_question` map)** — the
plan emits a `questions` array (3-5 entries, enforced by its
`output_schema`). The `map` node spawns one parallel branch per
question (`max_concurrency: 3`). Each branch is an isolated
`research_one_question` LLM invocation with web tools, instructed to
investigate exactly its assigned question. Outputs collect into
`question_findings` in input order, then `combine_findings` joins
them into a single `findings` Markdown document for downstream nodes.
`settings.max_concurrency: 4` is the graph-wide cap; the per-`map`
override (`max_concurrency: 3` on `research_each_question`) is
deliberately lower to leave headroom for the planner's tool calls
running alongside RAG.
## Local knowledge corpus
`knowledge_lookup` is a `rag` node — it runs hybrid (vector + keyword)
retrieval over every file in `knowledge/`. The directory ships with a
small `research-style-notes.md` so the RAG node has something to
retrieve against on a clean install; drop your own Markdown notes,
PDFs, or text files into `knowledge/` to bias the research toward
your local context.
The knowledge base is built once, at agent-load time, into
`~/.config/loki/agents/deep-research/knowledge_lookup.yaml`. Because
the node fully specifies its build config (`embedding_model`,
`chunk_size`, `chunk_overlap`), the build is non-interactive. Delete
that cached file after adding or changing knowledge to force a
rebuild.
## Sub-agent: report-writer
The `synthesize` node is an `agent` node that spawns the
`report-writer` sub-agent (`assets/agents/report-writer/`). This is
the agent-as-tool pattern: the orchestrating graph delegates the
writing phase to a focused sub-agent dedicated to coherent prose,
while the research phase uses different (typically cheaper) LLM nodes
for fast-and-many-question investigation.
The `report-writer` sub-agent has no tools — it cannot access the
web, cannot search, and cannot invent facts. It reads only the
findings it is given and produces a final Markdown report preserving
every inline citation. See `assets/agents/report-writer/README.md`
for details.
## Tools and tool scoping
This agent demonstrates Loki's three tool sources and how an `llm`
node's `tools:` whitelist scopes them per node.
The agent's full tool universe, declared in `graph.yaml`:
- **Global tools** (`global_tools`): `web_search_loki`,
`fetch_url_via_curl`, `search_arxiv` - Loki's built-in tool scripts.
- **MCP server** (`mcp_servers`): `ddg-search` - a DuckDuckGo web
search MCP server. Referenced in a whitelist as `mcp:ddg-search`.
- **Custom agent tool** (`tools.sh`): `classify_source` - a
deterministic source-credibility classifier shipped with this agent.
No node receives all of these. Each `llm` node's `tools:` whitelist
narrows the universe to exactly what that step needs:
| Node | `tools:` whitelist | Draws from |
|---|---|---|
| `plan`, `critique` | `[]` | nothing - pure reasoning |
| `research_one_question` | `web_search_loki`, `fetch_url_via_curl`, `search_arxiv`, `mcp:ddg-search` | global tools + MCP |
| `vet_sources` | `classify_source` | the custom tool only |
`research_one_question` (each parallel branch of the map) can search
and fetch but cannot classify sources; `vet_sources` can classify
sources but cannot touch the web. That separation is the point of the
`tools:` whitelist: a node gets only the tools its job calls for,
never the agent's full set.
The `classify_source` custom tool (`tools.sh`) takes a URL and returns
a credibility tier (government, academic, preprint, organization,
unverified) derived from the host and top-level domain. It is
deterministic - exactly the kind of logic a tool should own rather than
the LLM guessing.
Web search may require API-key configuration; see the
[Tools](https://github.com/Dark-Alex-17/loki/wiki/Tools) docs.
`fetch_url_via_curl`, `search_arxiv`, and `classify_source` work
without a key.
## Setup
`research_one_question` (each parallel branch of the `map`) uses the
`ddg-search` MCP server via `mcp:ddg-search`. It is one of Loki's
default MCP servers; make sure it is registered in
`~/.config/loki/mcp.json` (run `loki --install mcp_config` to restore
the default template if it is missing). If `ddg-search` is unavailable,
the branches still have their global web-search tools to fall back on.
The `synthesize` node spawns the `report-writer` sub-agent. Both
agents ship with `loki agents install`; if you install one manually,
install both so the agent reference resolves.
## Reflexion
The agent has two loops, both built with script nodes that route via
`_next`. The engine allows back-edges at runtime; the validator only
rejects cycles built from static `next` / `routes` edges, so script
`_next` loops are always allowed.
**Automated reflexion loop.** After the parallel research map and
`vet_sources`, the `critique` node reviews the merged findings
against the research plan and the source credibility assessment, and
emits `VERDICT: PASS` or `VERDICT: REVISE` with specific feedback.
`reflexion_gate.py` then:
- `PASS` -> continue to `synthesize`.
- `REVISE`, budget remaining -> loop back to `research_each_question`,
with the critique injected as `research_feedback` so every parallel
branch sees it on the retry.
- `REVISE`, budget spent -> continue to `synthesize` anyway (the human
approval step is the final backstop).
The budget is `MAX_REFLEXION_REVISIONS` in `reflexion_gate.py`
(default 2, so the research map runs at most 3 times per pass).
**Human-feedback loop.** At `approve` the user answers `accept`,
`reject`, or types their own feedback. A free-form answer routes via
the approval node's `on_other` to `incorporate_feedback.py`, which
folds that text into `research_feedback` and loops back to
`research_each_question` for another parallel pass.
`settings.max_loop_iterations` (40) is the engine's infinite-loop
backstop: it caps the total visits to any single node.
## Running
```sh
loki agents install # ships deep-research
loki -a deep-research "How does HTTP/3 differ from HTTP/2?"
loki -a deep-research "Recent advances in solid-state batteries"
loki -a deep-research # no prompt -> triggers ask_topic
```
## Anti-hallucination
- `research_one_question` (each map branch) is instructed to back
every claim with a real retrieved source and never to fabricate
URLs, titles, or DOIs.
- `vet_sources` classifies every cited source so weak sources are
visible to the critique step.
- `critique` independently reviews the merged findings and sends weak
or uncited work back for another parallel research pass.
- `synthesize` (the `report-writer` sub-agent) is grounded: it may use
only the gathered findings and must keep each claim's inline source.
It has no tools and cannot browse the web.
- `verify_sources` probes every cited URL / DOI with an HTTP HEAD
request and reports which are unreachable, so the human reviewer
sees broken citations before approving.
## Customizing
- **Loop budget.** `MAX_REFLEXION_REVISIONS` in `reflexion_gate.py`.
- **Map concurrency.** The `research_each_question` node's
`max_concurrency: 3` caps simultaneous web-research branches.
Raise to investigate more questions in parallel; lower to be gentle
on rate-limited providers.
- **Per-node model.** Add `model: anthropic:...` to any `llm` node.
Cheap models work well for `plan` / `critique` / `vet_sources`; the
heavy intelligence is needed in `research_one_question` and the
`report-writer` sub-agent.
- **Tool scope.** Narrow the `research_one_question` node's `tools:`
list to constrain where each branch looks (for example, drop
`web_search_loki` and `mcp:ddg-search` to force arXiv-only
research).
- **Local knowledge.** Drop files into `knowledge/` to bias every
research branch toward your local context (see the *Local
knowledge corpus* section above).
- **Different writer.** Replace `agent: report-writer` on the
`synthesize` node with the name of any other agent. The
orchestrator does not care what kind of agent the writer is.
- **Skip approval.** Point both `approve` routes at `end_accepted`,
or wire `verify_sources` straight to an `end` node.
## Files
```
assets/agents/deep-research/
graph.yaml - agent config + 17-node workflow
tools.sh - classify_source custom tool
README.md - this file
knowledge/
README.md - corpus-format notes
research-style-notes.md - starter knowledge file (replace with your notes)
scripts/
parse_request.py - _next: bootstrap_research, or ask_topic if no topic
bootstrap_research.py - fan-out source: next [plan, knowledge_lookup]
combine_findings.py - joins map output (question_findings) into findings
reflexion_gate.py - _next: research_each_question (revise) or synthesize
verify_sources.py - HTTP HEAD on cited URLs / DOIs
incorporate_feedback.py - _next: research_each_question, with user feedback
```
See also `assets/agents/report-writer/` — the sub-agent the
`synthesize` node spawns.
+293
View File
@@ -0,0 +1,293 @@
name: deep-research
description: |
Deep web research workflow. Plans an investigation, decomposes it
into sub-questions researched in parallel, grounds the work in a
local knowledge corpus, vets the credibility of cited sources, runs
a reflexion self-critique loop to revise weak or incomplete findings,
delegates the final write-up to a focused sub-agent, checks that the
cited sources are reachable, and gates the result behind human
approval. A reviewer's free-form feedback at the approval step feeds
back into another research pass.
This is the canonical Loki graph-agent reference: it exercises every
node type (script, llm, rag, map, agent, input, approval, end) and
both static fan-out and dynamic map fan-out.
version: "1.0"
temperature: 0.0
global_tools:
- web_search_loki.sh
- fetch_url_via_curl.sh
- search_arxiv.sh
mcp_servers:
- ddg-search
conversation_starters:
- "How does HTTP/3 differ from HTTP/2?"
- "Summarize recent advances in solid-state battery chemistry"
settings:
max_loop_iterations: 40
log_state_snapshots: false
validate_before_run: true
max_concurrency: 4
initial_state:
research_feedback: ""
research_attempts: 0
local_context: ""
local_sources: ""
start: parse_request
nodes:
parse_request:
id: parse_request
type: script
script: scripts/parse_request.py
next: bootstrap_research
ask_topic:
id: ask_topic
type: input
question: "What would you like me to research?"
validation: "len(input) > 0"
state_updates:
topic: "{{input}}"
next: bootstrap_research
bootstrap_research:
id: bootstrap_research
type: script
script: scripts/bootstrap_research.py
next: [plan, knowledge_lookup]
plan:
id: plan
type: llm
instructions: |
You are a research planner. Given a topic, produce a focused
research plan and decompose it into 3-5 specific sub-questions
that can each be researched independently in parallel.
The plan is a short narrative naming the key questions and the
kinds of sources that would be authoritative. The sub-questions
are precise, self-contained queries (each one is sent on its own
to a separate research worker, so they must be answerable
without each other's context).
prompt: "Research topic: {{topic}}"
tools: []
output_schema:
type: object
properties:
research_plan:
type: string
description: A short plan narrative.
questions:
type: array
items: { type: string }
minItems: 1
maxItems: 6
description: 3-5 specific, self-contained sub-questions.
required: [research_plan, questions]
next: research_each_question
knowledge_lookup:
id: knowledge_lookup
type: rag
documents:
- ./knowledge/
query: "{{topic}}"
top_k: 6
chunk_size: 1000
chunk_overlap: 100
state_updates:
local_context: "{{output.context}}"
local_sources: "{{output.sources}}"
next: research_each_question
research_each_question:
id: research_each_question
type: map
over: "{{questions}}"
as: question
branch: research_one_question
collect_into: question_findings
max_concurrency: 3
next: combine_findings
research_one_question:
id: research_one_question
type: llm
instructions: |
You are a web research assistant. Investigate the SINGLE question
given to you using your tools: search the web, fetch and read
pages, and search arXiv for academic sources.
Rules:
- Every factual claim must be backed by a real source you
actually retrieved. Never fabricate URLs, page titles,
authors, or DOIs.
- Prefer primary and authoritative sources over aggregators.
- Where sources disagree, report the disagreement rather than
papering over it.
- Put the URL (or DOI) inline next to each claim it supports.
Return organized findings in plain text. Do not include
meta-commentary about the process.
prompt: |
Research question: {{question}}
Local context that may help:
{{local_context}}
{{research_feedback}}
tools:
- web_search_loki
- fetch_url_via_curl
- search_arxiv
- mcp:ddg-search
max_iterations: 10
max_attempts: 2
temperature: 0.1
combine_findings:
id: combine_findings
type: script
script: scripts/combine_findings.py
next: vet_sources
vet_sources:
id: vet_sources
type: llm
instructions: |
You assess the credibility of the sources cited in a set of
research findings. For every distinct source URL in the findings,
call the `classify_source` tool to get its credibility tier. Then
summarize: which claims rest on HIGH-credibility sources, and
which rest on PREPRINT or UNVERIFIED sources and so need
corroboration. Do NOT do any new research -- assess only what is
already cited.
prompt: |
Findings to assess:
{{findings}}
tools:
- classify_source
max_iterations: 15
state_updates:
source_assessment: "{{output}}"
next: critique
critique:
id: critique
type: llm
instructions: |
You are a meticulous research reviewer. Judge whether the
findings below are good enough to synthesize a complete,
well-supported report that answers the research plan.
Mark the findings REVISE if ANY of these hold:
- A research-plan question is unanswered or only weakly
addressed.
- A factual claim has no source, or cites a source that looks
fabricated.
- The findings lean on a single source where corroboration is
needed.
- A key claim rests only on a PREPRINT or UNVERIFIED source,
per the source credibility assessment below.
- An obvious counter-perspective or recent development is
missing.
Otherwise mark them PASS.
Respond in EXACTLY this format, nothing else:
VERDICT: <PASS or REVISE>
FEEDBACK: <if REVISE, be specific and actionable -- name the gaps
and what kind of source would close them; if PASS, write "none">
prompt: |
Research plan:
{{research_plan}}
Findings under review:
{{findings}}
Source credibility assessment:
{{source_assessment}}
tools: []
state_updates:
critique: "{{output}}"
next: reflexion_gate
reflexion_gate:
id: reflexion_gate
type: script
script: scripts/reflexion_gate.py
next: synthesize
synthesize:
id: synthesize
type: agent
agent: report-writer
prompt: |
Research topic: {{topic}}
Findings (organized by sub-question, with inline citations):
{{findings}}
Source credibility assessment:
{{source_assessment}}
Produce the final report following your instructions.
timeout: 300
state_updates:
report: "{{output}}"
next: verify_sources
verify_sources:
id: verify_sources
type: script
script: scripts/verify_sources.py
next: approve
approve:
id: approve
type: approval
question: |
Research report on: {{topic}}
{{report}}
----
{{source_check}}
----
Accept this report? Pick "accept" or "reject", or type specific
feedback to send the research back for another pass.
options:
- "accept"
- "reject"
routes:
"accept": end_accepted
"reject": end_rejected
on_other: incorporate_feedback
state_updates:
decision: "{{choice}}"
incorporate_feedback:
id: incorporate_feedback
type: script
script: scripts/incorporate_feedback.py
end_accepted:
id: end_accepted
type: end
output: "{{report}}"
end_rejected:
id: end_rejected
type: end
output: "Research on '{{topic}}' was rejected and discarded."
@@ -0,0 +1,23 @@
# Local knowledge corpus for deep-research
The `knowledge_lookup` node in `graph.yaml` is a `rag` node that runs
hybrid (vector + keyword) retrieval over every file in this directory.
Drop your own notes, papers (PDFs), Markdown docs, or text files here
and they will be indexed into a per-agent knowledge base on first run.
Loki supports common file types out of the box: `.md`, `.txt`, `.pdf`,
`.html`, and others. Subdirectories are walked recursively.
A small starter file (`research-style-notes.md`) ships so the RAG
node has something non-empty to retrieve against on a clean install.
Replace or extend it with your own materials to bias the research
phase toward your local context.
To force the knowledge base to rebuild after you add or change files,
delete the cached index:
```sh
rm ~/.config/loki/agents/deep-research/knowledge_lookup.yaml
```
The next run will rebuild from the current contents of this directory.
@@ -0,0 +1,49 @@
# Research style notes
These are general principles the `deep-research` agent should keep in
mind regardless of topic. Replace this file with your own notes if you
want to bias retrieval toward your local context.
## What "good research" means here
- **Every factual claim cites a source you actually retrieved.** Never
fabricate URLs, page titles, authors, or DOIs.
- **Primary sources beat aggregators.** Prefer the original paper, the
RFC, the standards body, or the manufacturer over a blog summarizing
them.
- **Corroboration matters where stakes are high.** If a single source
makes a strong claim, look for a second independent source before
taking it as established.
- **Disagreement is information, not noise.** If two credible sources
disagree, report the disagreement and the reasoning on each side.
- **Old does not mean wrong.** A 2014 RFC is still authoritative if no
newer one has obsoleted it; check before assuming a source is stale.
## Source-tier heuristics
The `vet_sources` node uses these rough tiers to weigh credibility.
The custom tool `classify_source` (see `tools.sh`) implements this
deterministically by hostname / TLD.
- **HIGH:** government domains (`.gov`, `.mil`), academic institutions
(`.edu`, university subdomains), peer-reviewed journals, standards
bodies (IETF/RFCs, W3C, ISO, IEEE, NIST), and primary documents from
the entities being researched (e.g. a vendor's official spec page).
- **PREPRINT:** arXiv, bioRxiv, medRxiv, SSRN. Useful but not yet
peer-reviewed; treat numeric claims with extra caution.
- **ORGANIZATION:** established nonprofits, standards-adjacent groups,
industry consortia. Reliable for their stated mission but may have a
perspective.
- **UNVERIFIED:** general web pages, blogs, news aggregators, social
media. Useful for leads but should not be the only source for a
factual claim.
## Common pitfalls to flag in critique
- A claim cited only to a PREPRINT or UNVERIFIED source on a numeric
or contested point.
- A research-plan question that the findings address only obliquely.
- "Findings" that paraphrase a single source three times rather than
triangulating.
- Citation collisions where two sources are listed but turn out to
be the same study reported via different aggregators.
@@ -0,0 +1,18 @@
#!/usr/bin/env python3
"""Fan-out source for context loading.
Has no logic of its own. Exists so the static `next: [plan, knowledge_lookup]`
list on this node fans out into two parallel branches (the LLM planner and
the RAG knowledge lookup) as a single super-step. The validator requires
declared parallel-branch script outputs, so we emit an empty JSON object
explicitly here.
"""
import json
def main():
print(json.dumps({}))
if __name__ == "__main__":
main()
@@ -0,0 +1,39 @@
#!/usr/bin/env python3
"""Join the per-question map outputs into a single `findings` string.
The `research_each_question` map writes `question_findings` (an array,
one entry per sub-question, in input order). Downstream nodes
(`vet_sources`, `critique`, `synthesize`) read `{{findings}}` as a
single block, so this script renders the array as a Markdown document
with one section per question.
"""
import json
import os
def load_state():
path = os.environ.get("GRAPH_STATE_FILE")
if path:
with open(path) as f:
return json.load(f)
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
def main():
state = load_state()
questions = state.get("questions") or []
per_question = state.get("question_findings") or []
sections = []
for idx, q in enumerate(questions):
body = per_question[idx] if idx < len(per_question) else ""
if isinstance(body, dict) or isinstance(body, list):
body = json.dumps(body, indent=2)
sections.append(f"## {q}\n\n{body}")
findings = "\n\n".join(sections) if sections else "No findings gathered."
print(json.dumps({"findings": findings}))
if __name__ == "__main__":
main()
@@ -0,0 +1,41 @@
#!/usr/bin/env python3
"""Fold a reviewer's free-form feedback back into the research loop.
Runs when the user answers the approval step with their own text
instead of "accept" or "reject". That text (saved by the approval node
as `decision`) becomes `research_feedback`, and the graph loops back to
`research_each_question` for another informed pass (each sub-question is
re-researched in parallel with the new feedback in context). The
reflexion counter is reset so the user-driven pass gets a fresh revision
budget.
Routing (`_next`): always research_each_question.
"""
import json
import os
def load_state():
path = os.environ.get("GRAPH_STATE_FILE")
if path:
with open(path) as f:
return json.load(f)
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
def main():
state = load_state()
feedback = (state.get("decision") or "").strip()
output = {
"_next": "research_each_question",
"research_attempts": 0,
"research_feedback": (
"The user reviewed the report and asked for changes. Treat "
"this as the top priority for the next pass:\n\n" + feedback
),
}
print(json.dumps(output))
if __name__ == "__main__":
main()
@@ -0,0 +1,35 @@
#!/usr/bin/env python3
"""Entry router for deep-research.
Reads the caller's prompt from state. If it contains a usable research
topic, stores it as `topic` and falls through to the static `next`
(plan). If the prompt is empty, routes to `ask_topic` so the user can
supply one interactively.
Routing (`_next`):
- prompt present -> (no _next; static next: plan)
- prompt empty -> ask_topic
"""
import json
import os
def load_state():
path = os.environ.get("GRAPH_STATE_FILE")
if path:
with open(path) as f:
return json.load(f)
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
def main():
state = load_state()
prompt = (state.get("initial_prompt") or "").strip()
if prompt:
print(json.dumps({"topic": prompt}))
else:
print(json.dumps({"_next": "ask_topic"}))
if __name__ == "__main__":
main()
@@ -0,0 +1,76 @@
#!/usr/bin/env python3
"""Reflexion gate for deep-research.
Runs after `critique` has reviewed the current research findings. If the
critique's verdict is REVISE and the reflexion budget is not spent,
loops back to `research` with the critique attached as
`research_feedback`, so the retry is informed rather than a blind
re-run. Otherwise it proceeds to `synthesize`.
Routing (`_next`):
- verdict PASS -> synthesize
- verdict REVISE, budget remaining -> research_each_question (+ research_feedback)
- verdict REVISE, budget spent -> synthesize
Reflexion is a best-effort quality booster, not a hard gate: once the
budget is spent the workflow proceeds anyway, and the human approval
step is the final backstop.
"""
import json
import os
import re
# Automated revision passes allowed. `research` runs at most
# MAX_REFLEXION_REVISIONS + 1 times per user pass. Bump to allow more.
MAX_REFLEXION_REVISIONS = 2
def load_state():
path = os.environ.get("GRAPH_STATE_FILE")
if path:
with open(path) as f:
return json.load(f)
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
def as_int(value, default=0):
try:
return int(value)
except (TypeError, ValueError):
return default
def parse_verdict(critique):
"""Pull PASS/REVISE from the critique's `VERDICT:` line. Defaults to
PASS when no verdict line is found, so a malformed critique lets the
workflow proceed instead of burning the whole revision budget."""
match = re.search(r"VERDICT:\s*([A-Za-z]+)", critique, re.IGNORECASE)
if not match:
return "PASS"
return match.group(1).upper()
def main():
state = load_state()
critique = state.get("critique") or ""
verdict = parse_verdict(critique)
attempts = as_int(state.get("research_attempts"))
if verdict == "REVISE" and attempts < MAX_REFLEXION_REVISIONS:
feedback = (
"A reviewer judged the previous research pass incomplete. "
"Address every point in the critique below:\n\n" + critique
)
output = {
"_next": "research_each_question",
"research_attempts": attempts + 1,
"research_feedback": feedback,
}
else:
output = {"_next": "synthesize"}
print(json.dumps(output))
if __name__ == "__main__":
main()
@@ -0,0 +1,69 @@
#!/usr/bin/env python3
"""Check that the sources cited in the research report are reachable.
Scans the final report for URLs and DOIs, probes each with a HEAD
request, and writes a `source_check` summary into state so the human
reviewer sees broken citations at the approval step.
Times out per request so a slow source cannot stall the graph.
"""
import json
import os
import re
import urllib.error
import urllib.request
DOI_RE = re.compile(r"\b(10\.\d{4,9}/[-._;()/:A-Z0-9]+)", re.IGNORECASE)
URL_RE = re.compile(r"https?://[^\s)\]\}\"'>]+")
def load_state():
path = os.environ.get("GRAPH_STATE_FILE")
if path:
with open(path) as f:
return json.load(f)
return json.loads(os.environ.get("GRAPH_STATE", "{}"))
def reachable(url, timeout=5.0):
req = urllib.request.Request(url, method="HEAD")
try:
with urllib.request.urlopen(req, timeout=timeout) as resp:
return 200 <= resp.status < 400
except urllib.error.HTTPError as e:
return 200 <= e.code < 400
except Exception:
return False
def main():
state = load_state()
report = state.get("report") or ""
urls = sorted({u.rstrip(".,;)") for u in URL_RE.findall(report)})
dois = sorted(set(DOI_RE.findall(report)))
results = []
for url in urls:
ok = reachable(url)
results.append(f" {'OK' if ok else 'UNREACHABLE'} {url}")
for doi in dois:
url = f"https://doi.org/{doi}"
if url in urls:
continue
ok = reachable(url)
results.append(f" {'OK' if ok else 'UNREACHABLE'} DOI {doi} ({url})")
if not results:
summary = "No web sources were cited in the report."
else:
summary = (
f"Source reachability ({len(results)} checked):\n"
+ "\n".join(results)
)
print(json.dumps({"source_check": summary}))
if __name__ == "__main__":
main()
+39
View File
@@ -0,0 +1,39 @@
#!/usr/bin/env bash
set -e
# @env LLM_OUTPUT=/dev/stdout The output path
# @cmd Classify the credibility tier of a web source from its URL.
# A deterministic check based on the host and top-level domain. Use it
# to weigh how much trust to place in a source before relying on it.
# @option --url! The full source URL to classify
classify_source() {
# shellcheck disable=SC2154
local url="$argc_url"
local host="${url#*://}"
host="${host%%/*}"
host="${host##*@}"
host="${host%%:*}"
host="$(printf '%s' "$host" | tr '[:upper:]' '[:lower:]')"
local tier
case "$host" in
'')
tier="UNKNOWN - no host could be parsed from the URL" ;;
*.gov | *.gov.* | *.mil)
tier="HIGH - government source" ;;
*.edu | *.edu.* | *.ac.*)
tier="HIGH - academic institution" ;;
arxiv.org | *.arxiv.org | biorxiv.org | *.biorxiv.org | medrxiv.org | *.medrxiv.org | ssrn.com | *.ssrn.com)
tier="PREPRINT - not yet peer reviewed, corroborate before citing" ;;
wikipedia.org | *.wikipedia.org)
tier="TERTIARY - encyclopedia, good for orientation not citation" ;;
*.org | *.org.*)
tier="MEDIUM - organization site, check for institutional bias" ;;
*)
tier="UNVERIFIED - general web source, corroborate before citing" ;;
esac
printf '%s: %s\n' "${host:-<none>}" "$tier" >> "$LLM_OUTPUT"
}
-14
View File
@@ -1,14 +0,0 @@
# Jira AI Agent
## Overview
The Jira AI Agent is designed to assist with managing tasks within Jira projects, providing capabilities such as
creating, searching, updating, assigning, linking, and commenting on issues. Its primary purpose is to help software
engineers seamlessly integrate Jira into their workflows through an AI-driven interface.
## Configuration
This agent uses the official [Atlassian MCP Server](https://github.com/atlassian/atlassian-mcp-server). To use it,
ensure you have Node.js v18+ installed to run the local MCP proxy (`mcp-remote`).
The server uses OAuth 2.0 so it will automatically open your browser for you to sign in to your account. No manual
configuration is necessary!
-37
View File
@@ -1,37 +0,0 @@
name: Jira Agent
description: An AI agent that can assist with Jira tasks such as creating issues, searching for issues, and updating issues.
version: 0.1.0
agent_session: temp
mcp_servers:
- atlassian
instructions: |
You are a AI agent designed to assist with managing Jira tasks and helping software engineers utilize and integrate
Jira into their workflows. You can create, search, update, assign, link, and comment on issues in Jira.
## Create Issue (MANDATORY when creating a issue)
When a user prompts you to create a Jira issue:
1. Prompt the user for what Jira project they want the ticket created in
2. If the ticket type requires a parent issue:
a. Query Jira for potentially relevant parents
b. Prompt user for which parent to use, displaying the suggested list of parent issues
3. Create the issue with the following format:
```markdown
**Description:**
This section gives context and details about the issue.
**User Acceptance Criteria:**
# This section provides bullet points that function like a checklist of all the things that must be completed in
# order for the issue to be considered done.
* Example criteria one
* Example criteria two
```
4. Ask the user if the issue should be assigned to them
a. If yes, then assign the user to the newly created issue
Available tools:
{{__tools__}}
conversation_starters:
- What are the latest issues in my Jira project?
- Can you create a new Jira issue for me?
- What are my open Jira issues?
- Can you search for issues with the label "bug" in my Jira project?
+46
View File
@@ -0,0 +1,46 @@
# report-writer
A tiny, focused sub-agent that turns a set of research findings into a
single coherent final report. Reads only what it is given — does not
do independent research, does not access the web, does not invent
facts. It exists as a focused tool for orchestrating agents to
delegate the writing phase to.
## Why a separate agent?
This is an example of the **agent-as-tool** pattern in graph agents.
The `deep-research` graph agent's `synthesize` node is an `agent` node
that spawns this one (see `assets/agents/deep-research/graph.yaml`).
Separating the role has two practical benefits:
- The orchestrating agent can use a cheap model (or a high-temperature
exploratory one) for the research phase, while letting the writing
phase use a different (typically lower-temperature, possibly larger)
model dedicated to coherent prose.
- The writing prompt is owned by this agent's `config.yaml` rather
than buried inside another agent's graph. You can polish it
independently without touching the research flow.
## Standalone use
You can also use this agent directly if you have a set of findings you
want polished:
```sh
loki -a report-writer "Topic: X. Findings: <paste findings here>"
```
It will produce a single Markdown report following the rules in its
system prompt: executive summary at the top, grouped sections by
related sub-questions, every inline citation preserved verbatim, and a
final "Open questions / disagreements" section.
## What it will NOT do
- Search the web, fetch URLs, query an MCP server, or use any tool.
It has no tools configured.
- Invent facts beyond what is in the findings you give it.
- Strip or rewrite citations.
These constraints are the point of the agent existing: a writer that
the orchestrator can trust to stay in its lane.
+34
View File
@@ -0,0 +1,34 @@
name: report-writer
description: Polishes research findings into a clear, citation-preserving final report
version: 1.0.0
temperature: 0.2
instructions: |
You are a technical writer. You will be given:
- a research topic
- a set of findings, organized per sub-question, with inline
citations next to each claim
- a source-credibility assessment of the cited sources
Your job is to produce a single, well-organized final report:
Rules:
- Use ONLY the findings provided. Do not introduce facts from
your own memory. Do not speculate beyond what the findings
support.
- Preserve every inline citation. If a sentence in the findings
had a URL or DOI, the equivalent sentence in your report must
keep the same citation.
- Lead with a 2-3 sentence executive summary at the top.
- Organize the body so that related sub-questions are grouped,
not strictly one section per question. The findings are raw
material; the report should read as a single coherent answer
to the original topic.
- End with a short "Open questions / disagreements" section
naming anything the findings flagged as unresolved or
contested.
Output plain Markdown. No metadata, no JSON wrapper.
conversation_starters:
- "Polish these findings into a cited report"
+4 -5
View File
@@ -18,16 +18,15 @@ Sisyphus acts as the primary entry point, capable of handling complex tasks by c
- 🛠️ **Tool Integration**: Seamlessly uses system tools for building, testing, and file manipulation.
## Pro-Tip: Use an IDE MCP Server for Improved Performance
Many modern IDEs now include MCP servers that let LLMs perform operations within the IDE itself and use IDE tools. Using
an IDE's MCP server dramatically improves the performance of coding agents. So if you have an IDE, try adding that MCP
server to your config (see the [MCP Server docs](../../../docs/function-calling/MCP-SERVERS.md) to see how to configure
them), and modify the agent definition to look like this:
Many modern IDEs (JetBrains, VS Code, Cursor, Zed, etc.) expose MCP servers that let LLMs use IDE tools directly. Using
one dramatically improves the performance of coding agents. If you have one, add it to your loki config (see the
[MCP Server docs](../../../docs/function-calling/MCP-SERVERS.md)) and reference it in this agent's `mcp_servers:` list:
```yaml
# ...
mcp_servers:
- jetbrains
- your-ide-mcp-server
global_tools:
- fs_read.sh
+29 -12
View File
@@ -119,20 +119,21 @@ instructions: |
1. todo__init --goal "Add user profiles API endpoint"
2. todo__add --task "Explore existing API patterns"
3. todo__add --task "Implement profile endpoint"
4. todo__add --task "Verify with build/test"
5. agent__spawn --agent explore --prompt "Find existing API endpoint patterns, route structures, and controller conventions. Include code snippets."
6. agent__spawn --agent explore --prompt "Find existing data models and database query patterns. Include code snippets."
7. agent__collect --id <id1>
8. agent__collect --id <id2>
9. todo__done --id 1
10. agent__spawn --agent coder --prompt "<structured prompt using Coder Delegation Format above, including code snippets from explore results>"
11. agent__collect --id <coder_id>
12. todo__done --id 2
13. run_build
14. run_tests
15. todo__done --id 3
4. agent__spawn --agent explore --prompt "Find existing API endpoint patterns, route structures, and controller conventions. Include code snippets."
5. agent__spawn --agent explore --prompt "Find existing data models and database query patterns. Include code snippets."
6. agent__collect --id <id1>
7. agent__collect --id <id2>
8. todo__done --id 1
9. agent__spawn --agent coder --prompt "<structured prompt using Coder Delegation Format above, including code snippets from explore results>"
10. agent__collect --id <coder_id>
11. todo__done --id 2
```
Note: the `coder` agent is a graph agent that runs verification (build +
tests) and a bounded fix-loop internally. You do NOT need to spawn a
separate build/test step. A `CODER_COMPLETE` outcome means build and
tests already passed.
### Example 2: Architecture/design question (explore + oracle in parallel)
User: "How should I structure the authentication for this app?"
@@ -172,6 +173,22 @@ instructions: |
10. **Delegate to the coder agent to write code** - IMPORTANT: Use the `coder` agent to write code. Do not try to write code yourself except for trivial changes
11. **Always output a summary of changes when finished** - Make it clear to user's that you've completed your tasks
## Coder Outcomes
The `coder` agent is a graph agent that runs the implement -> verify_build
-> verify_tests -> fix_loop pipeline internally. It always returns one of
three sentinel outcomes:
- `CODER_COMPLETE` - implementation succeeded with build + tests green.
Continue with any follow-up todos.
- `CODER_REJECTED` - user rejected the plan at the approval gate (only
triggered for high-complexity plans). Do NOT re-spawn coder blindly;
ask the user what to change first.
- `CODER_FAILED` - the fix-loop exhausted its budget without producing
green build/tests. The failure output includes the last build and tests
output. Surface this to the user; consider spawning `oracle` for
diagnosis if the failure is unclear.
## When to Do It Yourself
- Simple command execution
+2 -2
View File
@@ -73,11 +73,11 @@ def to_args:
to_entries | .[] |
(.key | split("_") | join("-")) as $key |
if .value | type == "array" then
.value | .[] | "--\($key) \(. | escape_shell_word)"
.value | .[] | "--\($key)=\(. | escape_shell_word)"
elif .value | type == "boolean" then
if .value then "--\($key)" else "" end
else
"--\($key) \(.value | escape_shell_word)"
"--\($key)=\(.value | escape_shell_word)"
end;
[ to_args ] | join(" ")
EOF
+2 -2
View File
@@ -70,11 +70,11 @@ def to_args:
to_entries | .[] |
(.key | split("_") | join("-")) as $key |
if .value | type == "array" then
.value | .[] | "--\($key) \(. | escape_shell_word)"
.value | .[] | "--\($key)=\(. | escape_shell_word)"
elif .value | type == "boolean" then
if .value then "--\($key)" else "" end
else
"--\($key) \(.value | escape_shell_word)"
"--\($key)=\(.value | escape_shell_word)"
end;
[ to_args ] | join(" ")
EOF
+8
View File
@@ -0,0 +1,8 @@
---
enabled_mcp_servers: atlassian
---
You are the librarian for the company's Confluence and Jira knowledge bases. Your job is to help users find and retrieve
information from these platforms. Use all tools at your disposal to answer user queries.
Available Tools:
{{__tools__}}
+6 -4
View File
@@ -17,16 +17,18 @@ agent_session: null # Set a session to use when starting the agent.
name: <agent-name> # Name of the agent, used in the UI and logs
description: <description> # Description of the agent, used in the UI
version: 1 # Version of the agent
# Todo System & Auto-Continuation
# These settings help smaller models handle multi-step tasks more reliably.
# See docs/TODO-SYSTEM.md for detailed documentation.
# Auto-Continue (Todo System)
# The auto-continue system provides built-in task tracking for improved reliability.
# When enabled, the model can create todo lists and the system will automatically
# prompt it to continue when incomplete tasks remain.
# See the [Todo System documentation](https://github.com/Dark-Alex-17/loki/wiki/TODO-System) for more information
auto_continue: false # Enable automatic continuation when incomplete todos remain
max_auto_continues: 10 # Maximum number of automatic continuations before stopping
inject_todo_instructions: true # Inject the default todo tool usage instructions into the agent's system prompt
continuation_prompt: null # Custom prompt used when auto-continuing (optional; uses default if null)
# Sub-Agent Spawning System
# Enable this agent to spawn and manage child agents in parallel.
# See docs/AGENTS.md for detailed documentation.
# See https://github.com/Dark-Alex-17/loki/wiki/Agents for detailed documentation.
can_spawn_agents: false # Enable the agent to spawn child agents
max_concurrent_agents: 4 # Maximum number of agents that can run simultaneously
max_agent_depth: 3 # Maximum nesting depth for sub-agents (prevents runaway spawning)
+17 -8
View File
@@ -27,18 +27,18 @@ sync_models_url: > # URL to sync model changes from
https://raw.githubusercontent.com/Dark-Alex-17/loki/refs/heads/main/models.yaml
# ---- REPL Prompt ----
# Custom REPL left/right prompts; see the [REPL Prompt Documentation](./docs/REPL-PROMPT.md) for more information
# Custom REPL left/right prompts; see the [REPL Prompt Documentation](https://github.com/Dark-Alex-17/loki/wiki/REPL-Prompt) for more information
left_prompt:
'{color.red}{model}){color.green}{?session {?agent {agent}>}{session}{?role /}}{!session {?agent {agent}>}}{role}{?rag @{rag}}{color.cyan}{?session )}{!session >}{color.reset} '
right_prompt:
'{color.purple}{?session {?consume_tokens {consume_tokens}({consume_percent}%)}{!consume_tokens {consume_tokens}}}{color.reset}'
# ---- Vault ----
# See the [Vault documentation](./docs/VAULT.md) for more information on the Loki vault
# See the [Vault documentation](https://github.com/Dark-Alex-17/loki/wiki/Vault) for more information on the Loki vault
vault_password_file: null # Path to a file containing the password for the Loki vault (cannot be a secret template)
# ---- Function Calling ----
# See the [Tools documentation](./docs/function-calling/TOOLS.md) for more details
# See the [Tools documentation](https://github.com/Dark-Alex-17/loki/wiki/Tools) for more details
function_calling: true # Enables or disables function calling (Globally).
mapping_tools: # Alias for a tool or toolset
fs: 'fs_cat,fs_ls,fs_mkdir,fs_rm,fs_write,fs_read,fs_glob,fs_grep'
@@ -64,7 +64,6 @@ visible_tools: # Which tools are visible to be compiled (and a
# - get_current_weather.py
# - get_current_weather.ts
- get_current_weather.sh
- query_jira_issues.sh
# - search_arxiv.sh
# - search_wikipedia.sh
# - search_wolframalpha.sh
@@ -75,14 +74,24 @@ visible_tools: # Which tools are visible to be compiled (and a
# - web_search_tavily.sh
# ---- MCP Servers ----
# See the [MCP Servers documentation](./docs/MCP-SERVERS.md) for more details
# See the [MCP Servers documentation](https://github.com/Dark-Alex-17/loki/wiki/MCP-Servers) for more details
mcp_server_support: true # Enables or disables MCP servers (globally).
mapping_mcp_servers: # Alias for an MCP server or set of servers
git: github,gitmcp
enabled_mcp_servers: null # Which MCP servers to enable by default (e.g. 'github,slack,ddg-search')
# ---- Auto-Continue (Todo System) ----
# The auto-continue system provides built-in task tracking for improved reliability.
# When enabled, the model can create todo lists and the system will automatically
# prompt it to continue when incomplete tasks remain.
# See the [Todo System documentation](https://github.com/Dark-Alex-17/loki/wiki/TODO-System) for more information
auto_continue: false # Enable automatic continuation when incomplete todos remain (default: false)
max_auto_continues: 10 # Maximum number of automatic continuations before stopping (default: 10)
inject_todo_instructions: true # Inject default todo usage instructions into the system prompt (default: true)
continuation_prompt: null # Custom prompt used when auto-continuing. If null, uses built-in default
# ---- Session ----
# See the [Session documentation](./docs/SESSIONS.md) for more information
# See the [Session documentation](https://github.com/Dark-Alex-17/loki/wiki/Sessions) for more information
save_session: null # Controls the persistence of the session. If true, auto save; if false, don't auto-save save; if null, ask the user what to do
compression_threshold: 4000 # Compress the session when the token count reaches or exceeds this threshold
summarization_prompt: > # The text prompt used for creating a concise summary of session message
@@ -91,7 +100,7 @@ summary_context_prompt: > # The text prompt used for including the summar
'This is a summary of the chat history as a recap: '
# ---- RAG ----
# See the [RAG Docs](./docs/RAG.md) for more details.
# See the [RAG Docs](https://github.com/Dark-Alex-17/loki/wiki/RAG) for more details.
rag_embedding_model: null # Specifies the embedding model used for context retrieval
rag_reranker_model: null # Specifies the reranker model used for sorting retrieved documents; Loki uses Reciprocal Rank Fusion by default
rag_top_k: 5 # Specifies the number of documents to retrieve for answering queries
@@ -137,7 +146,7 @@ document_loaders:
sh -c "yek $1 --json | jq 'map({ path: .filename, contents: .content })'"
# ---- Clients ----
# See the [Clients documentation](./docs/clients/CLIENTS.md) for more details
# See the [Clients documentation](https://github.com/Dark-Alex-17/loki/wiki/Clients) for more details
clients:
# All clients have the following configuration:
# - type: xxxx
+14 -1
View File
@@ -1,5 +1,9 @@
---
# Everything in this section is optional
############################################
## Everything in this section is optional ##
############################################
# Role Configuration
name: <role-name> # The name of the role
model: openai:gpt-4o # The model to use for this role
temperature: 0.2 # The temperature to use for this role when querying the model
@@ -8,5 +12,14 @@ enabled_tools: fs_ls,fs_cat # A comma-separated list of tools to enabl
enabled_mcp_servers: github,gitmcp # A comma-separated list of MCP servers to enable for this role
prompt: null # A custom prompt to use for this role that will immediately query
# the model for output instead of using the instructions below
# Auto-Continue (Todo System)
# The auto-continue system provides built-in task tracking for improved reliability.
# When enabled, the model can create todo lists and the system will automatically
# prompt it to continue when incomplete tasks remain.
# See the [Todo System documentation](https://github.com/Dark-Alex-17/loki/wiki/TODO-System) for more information
auto_continue: false # Enable automatic continuation when incomplete todos remain (default: false)
max_auto_continues: 10 # Maximum number of automatic continuations before stopping (default: 10)
inject_todo_instructions: true # Inject default todo tool usage instructions into the system prompt (default: true)
continuation_prompt: null # Custom prompt used when auto-continuing. If null, uses built-in default
---
You are an expert at doing things. This is where you write the instructions for the role.
+427
View File
@@ -0,0 +1,427 @@
# Graph-based agent definition (full-featured reference)
# Location: <loki-config-dir>/agents/<agent-name>/graph.yaml
#
# A graph agent is defined by this file alone. An agent directory contains
# either a config.yaml (a normal LLM-loop agent) or a graph.yaml (a graph
# agent), never both. The presence of graph.yaml is what makes the agent
# a graph agent.
#
# This file is a reference: it documents every available field, themed
# around a deep web research workflow with parallel retrieval. It is not
# a runnable agent as-is. The `agent:`, `script:`, and `documents:` values
# point at things that would need to exist for a real agent. For a real,
# runnable deep-research graph agent, see assets/agents/deep-research/.
#
# Full documentation:
# https://github.com/Dark-Alex-17/loki/wiki/Graph-Agents
# ---------------------------------------------------------------------------
# Identity
# ---------------------------------------------------------------------------
name: deep-research-example # Agent name (should match the directory name)
description: | # Free-form prose describing the workflow
A reference workflow: triage a research request, retrieve local
context, branch on a script decision, run either a sub-agent or an
LLM research step, then gate the result behind human approval.
version: "1.0" # Graph schema version. Only "1.0" is accepted.
# ---------------------------------------------------------------------------
# Agent-level config (all optional)
# The same knobs a normal agent's config.yaml carries. In a graph agent they
# live here instead of in a config.yaml.
# ---------------------------------------------------------------------------
model: claude:claude-sonnet-4-6 # Default model for `llm` nodes that don't override it
temperature: 0.0 # Default sampling temperature for `llm` nodes
top_p: null # Default sampling top-p for `llm` nodes
global_tools: # Tool universe an `llm` node's `tools:` whitelist draws from
- web_search_loki.sh
- fetch_url_via_curl.sh
mcp_servers: # MCP servers an `llm` node may reference via `mcp:<server>`
- ddg-search
conversation_starters: # Suggested prompts surfaced in the UI
- "Research the current state of WebAssembly outside the browser"
# ---------------------------------------------------------------------------
# Agent variables (optional)
# Declared the same way as a normal agent's config.yaml `variables:` block.
# Each variable becomes available to:
# - LLM nodes via the template form `{{name}}` once seeded into state
# (see initial_state below).
# - Script nodes via the env var `LLM_AGENT_VAR_<UPPER_NAME>`.
# Values may be overridden at runtime with
# `loki -a <agent> --agent-variable <name> <value> "..."`.
# ---------------------------------------------------------------------------
variables:
- name: project_dir
description: |
Absolute path to the project directory.
default: "."
# ---------------------------------------------------------------------------
# Execution settings (all optional)
# ---------------------------------------------------------------------------
settings:
max_loop_iterations: 100 # Per-node visit cap. If one node id is entered more
# than this many times, execution aborts. Default 100.
timeout: 600 # Optional wall-clock cap (seconds) on the whole run,
# checked between node transitions.
log_state_snapshots: true # Log state before each node (debug/trace). Default true.
validate_before_run: true # Run the graph validator at startup. Default true.
max_concurrency: 4 # Cap on simultaneously running branches in any
# super-step (static fan-out OR a `map` node).
# Default 4. Per-`map` overrides this. See Parallel
# Execution below.
# ---------------------------------------------------------------------------
# Reducers (optional, required whenever two parallel branches write the same
# state key in the same super-step; otherwise the validator errors at load).
#
# A reducer says how two values for the same key get merged. Built-ins:
# append list += [value] (single value appended to a list)
# extend list += value (a list) (list-of-lists flattened by one level)
# concat "a\nb" (string join with newline separator)
# sum a + b (numeric add; ints stay ints)
# max max(a, b)
# min min(a, b)
# merge {**a, **b} (dict union, RHS wins on key collision)
# overwrite last-write-wins (explicit opt-in; B's value replaces A's)
#
# Keys not listed here have an implicit "single writer per super-step" rule:
# the validator rejects any graph where two parallel branches both write a
# key with no reducer.
# ---------------------------------------------------------------------------
reducers:
sources: append # The diamond below writes `sources` from both
# branches; append accumulates them into a list.
context: concat # Each branch contributes prose; concat joins them.
# ---------------------------------------------------------------------------
# Seed state (optional)
# Values placed into graph state before any node runs; reference anywhere via
# {{key}}.
#
# Note: `initial_prompt` is seeded automatically by Loki with the
# caller's prompt. So there's no need to set it here.
# ---------------------------------------------------------------------------
initial_state:
audience: "general reader"
# Seed an empty default for any key that a strict field (a node prompt /
# instructions / question / End output) references but that is only set on
# some paths. `refinement` is set only if the `refine` input node runs;
# seeding it "" keeps `finalize`'s strict prompt from failing on the
# approve-directly path.
refinement: ""
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
start: triage # ID of the first node to run (must exist in `nodes`)
# ---------------------------------------------------------------------------
# Nodes
# Each node is keyed by its id. The `id:` inside a node must match its key
# (it may also be omitted and thus Loki fills it in from the key).
#
# Node types: agent | script | approval | input | llm | rag | map | end
# ---------------------------------------------------------------------------
nodes:
# --- llm node -----------------------------------------------------------
# A one-shot LLM call (with an optional bounded tool-call loop). Runs in a
# fresh isolated context. Tools are strictly opt-in (see `tools`).
triage:
id: triage
type: llm
description: Classify the research request and extract its topic.
instructions: | # Optional system prompt (templated against state)
You triage research requests for a {{audience}} audience.
prompt: | # Required user prompt (templated against state)
Classify this request and extract the core research topic:
{{initial_prompt}}
tools: [] # Tool whitelist. Omitted or [] = no tools at all.
# A list narrows to exactly those entries.
output_schema: # Optional JSON Schema. The output is parsed to JSON
type: object # and its top-level object keys auto-merge into state
properties: # (so `topic` / `needs_deep_dive` become {{topic}} etc).
topic: { type: string }
needs_deep_dive: { type: boolean }
required: [topic, needs_deep_dive]
state_updates: # {{output}} = this node's result (here, the parsed object)
triage_result: "{{output}}"
# --- Polymorphic `next` -----------------------------------------------
# A single string runs the next node sequentially (e.g. `next: retrieve`).
# A list runs all listed nodes in parallel as one BSP super-step
# (for more info on BSP, see https://en.wikipedia.org/wiki/Bulk_synchronous_parallel).
# Their writes are merged via `reducers:` at the join. Branches converge
# implicitly when they all route to the same downstream node (here,
# `synthesize`). See the diamond:
#
# triage
# / \
# retrieve web_search (run concurrently)
# \ /
# synthesize (join; fires once after both finish)
next: [retrieve, web_search]
# --- rag node (parallel branch 1 of the diamond) ------------------------
# Hybrid (vector + keyword) retrieval against a per-node knowledge base.
# The knowledge base is built once, at agent load time, into
# <agent-dir>/retrieve.yaml (named after this node's id).
retrieve:
id: retrieve
type: rag
documents: # Required. Files, directories, URLs, loader paths.
- ./knowledge/ # relative paths resolve against the agent directory
- https://example.com/reference
query: "{{topic}}" # Retrieval query (templated). Default: {{initial_prompt}}.
top_k: 5 # Chunks to retrieve. Default = the KB's own top_k.
timeout: 120 # Retrieval timeout in seconds. Default 120.
# Knowledge-base build config (optional; used only when the KB is first
# built). When embedding_model + chunk_size + chunk_overlap are all set,
# the KB builds with no interactive prompts (works in non-interactive runs).
embedding_model: openai:text-embedding-3-small
chunk_size: 1000
chunk_overlap: 100
reranker_model: null # Optional reranker for hybrid-search results
batch_size: 100 # Optional embedding-request batch size
state_updates: # {{output}} = { context: <str>, sources: [<path>, ...] }
context: "{{output.context}}" # writes `context` -> `reducers.context = concat`
sources: "{{output.sources}}" # writes `sources` -> `reducers.sources = append`
next: synthesize # Joins with web_search at `synthesize`.
# --- llm node (parallel branch 2 of the diamond) ------------------------
# Runs concurrently with `retrieve`. Both branches write `context` and
# `sources`; the validator confirms both keys have a reducer declared, and
# the BSP scheduler merges them at the join.
web_search:
id: web_search
type: llm
instructions: "You are a web researcher. Cite every claim."
prompt: "Web research: {{topic}}. Return findings and sources."
tools:
- web_search_loki
- mcp:ddg-search
output_schema:
type: object
properties:
context: { type: string }
sources:
type: array
items: { type: string }
required: [context, sources]
# When `output_schema` is set, top-level keys auto-merge into state, so
# `context` and `sources` are produced without needing `state_updates`.
next: synthesize # Joins with retrieve at `synthesize`.
# --- script node (the diamond's join; also dispatches) -----------------
# Runs a .sh / .py / .ts script. The script receives state via the
# GRAPH_STATE env var (inline JSON) or GRAPH_STATE_FILE (path to a JSON
# file, used when state exceeds 32 KiB). Exactly one is set. It must print
# a single JSON object on stdout: keys merge into state, and the reserved
# `_next` key (if present) overrides routing.
#
# The script also receives these env vars (parity with bash tools called
# from normal agents):
# GRAPH_STATE / GRAPH_STATE_FILE state payload (one of the two is set)
# LLM_ROOT_DIR loki config dir
# LLM_PROMPT_UTILS_FILE path to .shared/prompt-utils.sh
# LLM_AGENT_DATA_DIR this agent's data directory
# LLM_AGENT_VAR_<NAME> one per declared `variables:` entry
# PATH with loki's functions bin dir prepended
# CLICOLOR_FORCE / FORCE_COLOR so child tools emit ANSI colors
# The script's working directory is loki's invocation CWD (not the agent
# directory), matching the behavior of bash tools.
#
# This node fires once: after both `retrieve` and `web_search` finish.
# The BSP scheduler dedups the two incoming edges into a single frontier
# entry, applies the staged branch writes through the reducers, then runs
# this node against the merged state. Inside the script, `context` is the
# concatenated text of both branches and `sources` is the combined list.
synthesize:
id: synthesize
type: script
script: scripts/synthesize.py # Path relative to the agent directory
timeout: 30 # Seconds. Default 30.
state_updates: # Applied after the stdout JSON is merged
decided_for: "{{topic}}"
next: summarize # Default route if the script emits no `_next`
fallback: summarize # Route taken if the script fails (crash / bad JSON)
# This script is expected to emit `_next: deep_dive` (or `_next: subjects_map`
# to demonstrate the map node below), or no `_next` (then `next` is used).
# Targets reached only via the script's dynamic `_next` get an
# "unreachable" warning from the validator. This is expected for `_next`-routed
# targets.
# --- agent node ---------------------------------------------------------
# Spawns a full Loki sub-agent and waits for it. The child uses its own
# tool stack. Agent nodes have no `tools:` field. No schema hint is
# injected even when `output_schema` is set (unlike llm nodes).
deep_dive:
id: deep_dive
type: agent
agent: deep-research # Name of an existing Loki agent to spawn
prompt: | # User message sent to the child (templated)
Research {{topic}} in depth. Existing context:
{{context}}
timeout: 600 # Optional wall-clock cap, seconds. Default 300.
output_schema: # Optional. Same extraction as llm nodes
type: object
properties:
summary: { type: string }
findings:
type: array
items: { type: string }
required: [summary, findings]
state_updates:
research: "{{output}}"
next: review # Required for agent nodes
# --- map node (Dynamic fan-out. Think: LangGraph's `Send` API) ----------------
# Spawns one parallel sub-branch per item in `over`. Each sub-branch runs
# the node referenced by `branch:` with the item bound to `as:`. Outputs
# collect into the array named by `collect_into:`, preserving input order.
#
# Reach via `synthesize`'s `_next: subjects_map`. The producer is expected
# to have written a list at `subjects` (e.g. an upstream LLM node with an
# `output_schema` returning {"subjects": ["a", "b", "c"]}).
subjects_map:
id: subjects_map
type: map
over: "{{subjects}}" # Required. List expression resolved from state.
# Empty list is allowed. It means no branches spawn,
# and thus `collect_into` is written as [].
as: subject # Required. Per-branch state key holding the
# current item. Read with {{subject}} inside
# the branch node's prompt.
branch: research_subject # Required. Node id to invoke per item.
# Must point to an llm | agent | rag | script
# node satisfying the map branch contract:
# - no `next:` (atomic, joined at map exit)
# - no `state_updates:` other than via the
# map's `collect_into` channel
# - no `output_schema:` (top-level merge
# would clash with collect_into)
# Validator enforces all three.
collect_into: subject_findings # Required. State key for the array of
# per-branch outputs, in input order
# (not spawn-finish order).
max_concurrency: 3 # Optional per-map cap. Defaults to
# settings.max_concurrency above.
output_key: output # Optional. State key the branch's output
# appears under. Default "output". Useful
# only if the branch reads its own bound
# name back (rare).
next: aggregate_subjects # Where to go after all sub-branches finish.
# Branch node for subjects_map. Each invocation receives a different
# `subject` in state. The branch is "atomic", meaning it cannot route on
# its own; the surrounding `map` joins after all invocations finish.
research_subject:
id: research_subject
type: llm
instructions: "Research one subject deeply for a {{audience}} audience."
prompt: "Research {{subject}}: pull the key facts and one citation."
tools:
- web_search_loki
# No `next:`, `state_updates:`, or `output_schema:` here. Map branches
# have a strict contract (see `subjects_map.branch` comment).
# Aggregator that runs after the map joins. Reads the collected list.
aggregate_subjects:
id: aggregate_subjects
type: llm
instructions: "Combine N per-subject reports into one cohesive summary."
prompt: |
Per-subject reports (in original input order):
{{subject_findings}}
state_updates:
research: "{{output}}"
next: review
# --- llm node with a narrowed tool whitelist ----------------------------
summarize:
id: summarize
type: llm
instructions: "You write concise research summaries for a {{audience}} audience."
prompt: "Summarize the topic {{topic}}, using your tools as needed."
tools: # Narrow whitelist: exactly these entries, nothing else
- web_search_loki # an exact global-tool / custom-tool name
- mcp:ddg-search # `mcp:<server>` includes that server's functions
model: claude:claude-haiku-4-5 # Optional per-node model override
temperature: 0.3 # Optional per-node sampling override
max_attempts: 2 # Retry count on transient errors only. Default 1.
max_iterations: 10 # Tool-call-loop turn cap. Default 10.
fallback: review # Route here if all attempts fail
timeout: 300 # Optional node wall-clock cap, seconds (unset = no timeout)
state_updates:
research: "{{output}}"
next: review # Required for llm nodes: the success route
# --- approval node ------------------------------------------------------
# Human-in-the-loop checkpoint. `user__ask` always offers a free-form
# "type your own answer" option, so `on_other` is required.
review:
id: review
type: approval
question: |
Proposed research result for {{topic}}:
{{research}}
Approve?
options: # The listed choices shown to the user
- "yes"
- "no"
routes: # Map each listed option to its next node
"yes": finalize
"no": rejected_end
on_other: refine # Required: route for ANY answer not in `routes`
state_updates:
decision: "{{choice}}" # {{choice}} = the chosen option or the free-form text
# --- input node ---------------------------------------------------------
# Collects a free-form string from the user.
refine:
id: refine
type: input
question: "What should be changed about the research result?"
default: "tighten the summary" # Optional: used if the user submits empty input.
# Note: a substituted default is not re-validated,
# so make sure it would satisfy `validation`.
validation: "len(input) > 0" # Optional length predicate: len(input) <op> N,
# <op> in > >= < <= == . Length only -- no regex.
state_updates:
refinement: "{{input}}" # {{input}} = the user's text
next: finalize # Required for input nodes: the success route
# --- llm node (final synthesis) -----------------------------------------
finalize:
id: finalize
type: llm
prompt: |
Produce the final research report for {{topic}}.
Result so far: {{research}}
Requested refinement (if any): {{refinement}}
state_updates:
final_report: "{{output}}"
next: done
# --- end nodes ----------------------------------------------------------
# Terminate the graph. `output` (templated, lenient interpolation) becomes
# the graph's final result. A graph needs at least one `end` node.
done:
id: done
type: end
state_updates: # Optional: applied before `output` is rendered
status: "completed"
output: |
[{{status}}] {{final_report}}
Sources: {{sources}}
rejected_end:
id: rejected_end
type: end
output: "Research on {{topic}} was not approved."
-8
View File
@@ -487,14 +487,6 @@
thinking:
type: enabled
budget_tokens: 16000
- name: claude-3-5-haiku-20241022
max_input_tokens: 200000
max_output_tokens: 8192
require_max_tokens: true
input_price: 0.8
output_price: 4
supports_vision: true
supports_function_calling: true
# Links:
# - https://docs.mistral.ai/getting-started/models/models_overview/
+29
View File
@@ -1,9 +1,11 @@
use crate::client::{ModelType, list_models};
use crate::config::paths;
use crate::config::{AppConfig, Config, list_agents, list_sessions};
use crate::utils::list_file_names;
use crate::vault::Vault;
use clap_complete::{CompletionCandidate, Shell, generate};
use clap_complete_nushell::Nushell;
use std::env;
use std::ffi::OsStr;
use std::io;
@@ -94,9 +96,36 @@ pub(super) fn macro_completer(current: &OsStr) -> Vec<CompletionCandidate> {
.collect()
}
fn extract_agent_from_args() -> Option<String> {
let args: Vec<String> = env::args().collect();
let mut i = 0;
while i < args.len() {
let arg = &args[i];
if let Some(value) = arg.strip_prefix("--agent=") {
return Some(value.to_string());
}
if (arg == "--agent" || arg == "-a") && i + 1 < args.len() {
return Some(args[i + 1].clone());
}
i += 1;
}
None
}
pub(super) fn session_completer(current: &OsStr) -> Vec<CompletionCandidate> {
let cur = current.to_string_lossy();
let sessions = if let Some(agent_name) = extract_agent_from_args() {
let sessions_dir = paths::agent_data_dir(&agent_name).join("sessions");
list_file_names(sessions_dir, ".yaml")
} else {
list_sessions()
};
sessions
.into_iter()
.filter(|s| s.starts_with(&*cur))
.map(CompletionCandidate::new)
+48 -2
View File
@@ -4,9 +4,10 @@ use crate::cli::completer::{
ShellCompletion, agent_completer, macro_completer, model_completer, rag_completer,
role_completer, secrets_completer, session_completer,
};
use crate::config::{AssetCategory, InstallFilter};
use anyhow::{Context, Result};
use clap::ValueHint;
use clap::{Parser, crate_authors, crate_description, crate_name, crate_version};
use clap::{Parser, crate_authors, crate_description, crate_version};
use clap_complete::ArgValueCompleter;
use is_terminal::IsTerminal;
use std::io::{Read, stdin};
@@ -14,7 +15,7 @@ use std::io::{Read, stdin};
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
#[command(
name = crate_name!(),
name = "loki",
author = crate_authors!(),
version = crate_version!(),
about = crate_description!(),
@@ -82,6 +83,18 @@ pub struct Cli {
/// Build all configured Bash tool scripts
#[arg(long)]
pub build_tools: bool,
/// Reinstall bundled assets, overwriting any local changes
#[arg(long, value_name = "CATEGORY", value_enum)]
pub install: Option<AssetCategory>,
/// Install assets from a remote git repository (URL may be suffixed with #<ref>)
#[arg(long, value_name = "GIT_URL")]
pub install_from: Option<String>,
/// Restrict --install-from to a single asset category
#[arg(long, value_name = "CATEGORY", value_enum, requires = "install_from")]
pub filter: Option<InstallFilter>,
/// Overwrite all conflicts without prompting (used with --install-from)
#[arg(long, requires = "install_from")]
pub install_force: bool,
/// Sync models updates
#[arg(long)]
pub sync_models: bool,
@@ -133,6 +146,12 @@ pub struct Cli {
/// Generate static shell completion scripts
#[arg(long, value_name = "SHELL", value_enum)]
pub completions: Option<ShellCompletion>,
/// Update Loki to the latest release, or to a specific version
#[arg(long, value_name = "VERSION")]
pub update: Option<Option<String>>,
/// With --update, update even if Loki was installed via a package manager
#[arg(long, requires = "update")]
pub force: bool,
}
impl Cli {
@@ -392,4 +411,31 @@ mod tests {
let cli = parse(&["--macro", "my-macro"]);
assert_eq!(cli.macro_name, Some("my-macro".to_string()));
}
#[test]
fn parse_update_flag_no_value() {
let cli = parse(&["--update"]);
assert_eq!(cli.update, Some(None));
}
#[test]
fn parse_update_flag_with_version() {
let cli = parse(&["--update", "v0.4.0"]);
assert_eq!(cli.update, Some(Some("v0.4.0".to_string())));
}
#[test]
fn parse_update_with_force() {
let cli = parse(&["--update", "--force"]);
assert_eq!(cli.update, Some(None));
assert!(cli.force);
}
#[test]
fn parse_force_without_update_fails() {
assert!(Cli::try_parse_from(["loki", "--force"]).is_err());
}
}
+29 -24
View File
@@ -117,33 +117,38 @@ async fn prepare_chat_completions(
/// So this function injects the Claude Code system prompt into the request
/// body to make it a valid request.
fn inject_oauth_system_prompt(body: &mut Value) {
let prefix_block = json!({
"type": "text",
"text": CLAUDE_CODE_PREFIX,
});
match body.get("system") {
Some(Value::String(existing)) => {
let existing_block = json!({
"type": "text",
"text": existing,
});
body["system"] = json!([prefix_block, existing_block]);
let existing_text = match body.get("system") {
Some(Value::String(s)) => {
if s.starts_with(CLAUDE_CODE_PREFIX) {
return;
}
Some(Value::Array(_)) => {
if let Some(arr) = body["system"].as_array_mut() {
let already_injected = arr
(!s.is_empty()).then(|| s.clone())
}
Some(Value::Array(blocks)) => {
let already_injected = blocks.iter().any(|b| {
b.get("text")
.and_then(|t| t.as_str())
.map(|t| t.starts_with(CLAUDE_CODE_PREFIX))
.unwrap_or(false)
});
if already_injected {
return;
}
let joined: Vec<String> = blocks
.iter()
.any(|block| block["text"].as_str() == Some(CLAUDE_CODE_PREFIX));
if !already_injected {
arr.insert(0, prefix_block);
}
}
}
_ => {
body["system"] = json!([prefix_block]);
}
.filter_map(|b| b.get("text").and_then(|t| t.as_str()).map(String::from))
.collect();
(!joined.is_empty()).then(|| joined.join("\n\n"))
}
_ => None,
};
let merged = match existing_text {
Some(rest) => format!("{}\n\n{}", CLAUDE_CODE_PREFIX, rest),
None => CLAUDE_CODE_PREFIX.to_string(),
};
body["system"] = json!([{ "type": "text", "text": merged }]);
}
pub async fn claude_chat_completions(
+8 -3
View File
@@ -1,6 +1,6 @@
use super::*;
use crate::config::paths;
use crate::config::{RenderMode, paths};
use crate::{
config::{AppConfig, Input, RequestContext},
function::{FunctionDeclaration, ToolCall, ToolResult, eval_tool_calls},
@@ -418,7 +418,8 @@ pub async fn call_chat_completions(
abort_signal: AbortSignal,
) -> Result<(String, Vec<ToolResult>)> {
let is_child_agent = ctx.current_depth > 0;
let spinner_message = if is_child_agent { "" } else { "Generating" };
let suppress_spinner = is_child_agent || ctx.render_mode == RenderMode::Silent;
let spinner_message = if suppress_spinner { "" } else { "Generating" };
let ret = abortable_run_with_spinner(
client.chat_completions(input.clone()),
spinner_message,
@@ -459,10 +460,14 @@ pub async fn call_chat_completions_streaming(
) -> Result<(String, Vec<ToolResult>)> {
let (tx, rx) = unbounded_channel();
let mut handler = SseHandler::new(tx, abort_signal.clone());
let silent = ctx.render_mode == RenderMode::Silent;
if silent {
handler.set_silent(true);
}
let (send_ret, render_ret) = tokio::join!(
client.chat_completions_streaming(input, &mut handler),
render_stream(rx, client.app_config(), abort_signal.clone()),
render_stream(rx, client.app_config(), abort_signal.clone(), silent),
);
if handler.abort().aborted() {
+10 -5
View File
@@ -94,21 +94,21 @@ impl MessageContent {
match self {
MessageContent::Text(text) => multiline_text(text),
MessageContent::Array(list) => {
let (mut concated_text, mut files) = (String::new(), vec![]);
let (mut concatenated_text, mut files) = (String::new(), vec![]);
for item in list {
match item {
MessageContentPart::Text { text } => {
concated_text = format!("{concated_text} {text}")
concatenated_text = format!("{concatenated_text} {text}")
}
MessageContentPart::ImageUrl { image_url } => {
files.push(resolve_url_fn(&image_url.url))
}
}
}
if !concated_text.is_empty() {
concated_text = format!(" -- {}", multiline_text(&concated_text))
if !concatenated_text.is_empty() {
concatenated_text = format!(" -- {}", multiline_text(&concatenated_text))
}
format!(".file {}{}", files.join(" "), concated_text)
format!(".file {}{}", files.join(" "), concatenated_text)
}
MessageContent::ToolCalls(MessageContentToolCalls {
tool_results, text, ..
@@ -227,9 +227,14 @@ pub fn patch_messages(messages: &mut Vec<Message>, model: &Model) {
}
pub fn extract_system_message(messages: &mut Vec<Message>) -> Option<String> {
if messages.is_empty() {
return None;
}
if messages[0].role.is_system() {
let system_message = messages.remove(0);
return Some(system_message.content.to_text());
}
None
}
+47 -27
View File
@@ -2,9 +2,9 @@ use super::{ToolCall, catch_error};
use crate::utils::AbortSignal;
use anyhow::{Context, Result, anyhow, bail};
use eventsource_stream::Eventsource;
use futures_util::{Stream, StreamExt};
use reqwest::RequestBuilder;
use reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt};
use reqwest::{RequestBuilder, header};
use serde_json::Value;
use tokio::sync::mpsc::UnboundedSender;
@@ -16,6 +16,7 @@ pub struct SseHandler {
last_tool_calls: Vec<ToolCall>,
max_call_repeats: usize,
call_repeat_chain_len: usize,
silent: bool,
}
impl SseHandler {
@@ -28,14 +29,24 @@ impl SseHandler {
last_tool_calls: Vec::new(),
max_call_repeats: 2,
call_repeat_chain_len: 3,
silent: false,
}
}
pub fn set_silent(&mut self, silent: bool) {
self.silent = silent;
}
pub fn text(&mut self, text: &str) -> Result<()> {
if text.is_empty() {
return Ok(());
}
self.buffer.push_str(text);
if self.silent {
return Ok(());
}
let ret = self
.sender
.send(SseEvent::Text(text.to_string()))
@@ -193,23 +204,13 @@ pub async fn sse_stream<F>(builder: RequestBuilder, mut handle: F) -> Result<()>
where
F: FnMut(SseMessage) -> Result<bool>,
{
let mut es = builder.eventsource()?;
while let Some(event) = es.next().await {
match event {
Ok(Event::Open) => {}
Ok(Event::Message(message)) => {
let message = SseMessage {
event: message.event,
data: message.data,
};
if handle(message)? {
break;
}
}
Err(err) => {
match err {
EventSourceError::StreamEnded => {}
EventSourceError::InvalidStatusCode(status, res) => {
let res = builder
.header(header::ACCEPT, "text/event-stream")
.header(header::CACHE_CONTROL, "no-store")
.send()
.await?;
let status = res.status();
if !status.is_success() {
let text = res.text().await?;
let data: Value = match text.parse() {
Ok(data) => data,
@@ -221,19 +222,38 @@ where
}
};
catch_error(&data, status.as_u16())?;
return Ok(());
}
EventSourceError::InvalidContentType(header_value, res) => {
let content_type = res
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(|value| value.to_string());
let is_event_stream = content_type
.as_deref()
.map(|ct| ct.starts_with("text/event-stream"))
.unwrap_or(false);
if !is_event_stream {
let header_value = content_type.unwrap_or_default();
let text = res.text().await?;
bail!(
"Invalid response event-stream. content-type: {}, data: {text}",
header_value.to_str().unwrap_or_default()
);
bail!("Invalid response event-stream. content-type: {header_value}, data: {text}");
}
_ => {
bail!("{}", err);
let mut es = res.bytes_stream().boxed().eventsource();
while let Some(event) = es.next().await {
match event {
Ok(message) => {
let message = SseMessage {
event: message.event,
data: message.data,
};
if handle(message)? {
break;
}
}
es.close();
Err(err) => {
bail!("{err}");
}
}
}
+301 -54
View File
@@ -11,6 +11,8 @@ use crate::config::prompts::{
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
};
use crate::graph::{Graph, GraphParser, NodeType};
use crate::rag::RagInitConfig;
use crate::vault::SECRET_RE;
use anyhow::{Context, Result};
use fancy_regex::Captures;
@@ -37,12 +39,13 @@ pub struct Agent {
session_dynamic_instructions: Option<String>,
functions: Functions,
rag: Option<Arc<Rag>>,
graph_rags: HashMap<String, Arc<Rag>>,
model: Model,
vault: GlobalVault,
}
impl Agent {
pub fn install_builtin_agents() -> Result<()> {
pub fn install_builtin_agents(force: bool) -> Result<()> {
info!(
"Installing built-in agents in {}",
paths::agents_data_dir().display()
@@ -62,7 +65,7 @@ impl Agent {
#[cfg_attr(not(unix), expect(unused))]
let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py"));
if file_path.exists() {
if file_path.exists() && !force {
debug!(
"Agent file already exists, skipping: {}",
file_path.display()
@@ -97,10 +100,28 @@ impl Agent {
let loaders = app.document_loaders.clone();
let rag_path = paths::agent_rag_file(name, DEFAULT_AGENT_NAME);
let config_path = paths::agent_config_file(name);
let mut agent_config = if config_path.exists() {
AgentConfig::load(&config_path)?
} else {
bail!("Agent config file not found at '{}'", config_path.display())
let graph_path = paths::agent_graph_file(name);
let mut graph_for_rag: Option<Graph> = None;
let mut agent_config = match (config_path.exists(), graph_path.exists()) {
(true, true) => bail!(
"Agent '{name}' has both config.yaml and graph.yaml. A graph agent \
is defined by graph.yaml alone; a normal agent by config.yaml alone. \
Remove one of the two files."
),
(true, false) => AgentConfig::load(&config_path)?,
(false, true) => {
let parser = GraphParser::new(&agent_data_dir);
let graph = parser
.load_from_file(&graph_path)
.with_context(|| format!("Failed to load graph.yaml for agent '{name}'"))?;
let config = AgentConfig::from_graph(name, &graph);
graph_for_rag = Some(graph);
config
}
(false, false) => bail!(
"Agent '{name}' has neither a config.yaml nor a graph.yaml at '{}'",
agent_data_dir.display()
),
};
let mut functions = Functions::init_agent(name, &agent_config.global_tools)?;
@@ -138,44 +159,16 @@ impl Agent {
.prompt()?;
}
if ans {
let mut document_paths = vec![];
for path in &agent_config.documents {
if is_url(path) {
document_paths.push(path.to_string());
} else if is_loader_protocol(&loaders, path) {
let (protocol, document_path) = path
.split_once(':')
.with_context(|| "Invalid loader protocol path")?;
let resolved_path = resolve_home_dir(document_path);
let new_path = if Path::new(&resolved_path).is_relative() {
safe_join_path(&agent_data_dir, resolved_path)
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?
} else {
PathBuf::from(&resolved_path)
};
document_paths.push(format!("{}:{}", protocol, new_path.display()));
} else if Path::new(&resolve_home_dir(path)).is_relative() {
let new_path = safe_join_path(&agent_data_dir, path)
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?;
document_paths.push(new_path.display().to_string())
} else {
document_paths.push(path.to_string())
}
}
let document_paths =
resolve_document_paths(&agent_config.documents, &loaders, &agent_data_dir)?;
let key = RagKey::Agent(name.to_string());
let app_clone = app.clone();
let rag_path_clone = rag_path.clone();
let abort = abort_signal.clone();
let rag = app_state
.rag_cache
.load_with(key, || async move {
Rag::init(
&app_clone,
"rag",
&rag_path_clone,
&document_paths,
abort_signal,
)
.await
Rag::init(&app_clone, "rag", &rag_path_clone, &document_paths, abort).await
})
.await?;
Some(rag)
@@ -186,6 +179,23 @@ impl Agent {
None
};
let graph_rags = match &graph_for_rag {
Some(graph) => {
init_graph_rags(
app,
app_state,
name,
graph,
&agent_data_dir,
&loaders,
info_flag,
abort_signal.clone(),
)
.await?
}
None => HashMap::new(),
};
if agent_config.auto_continue {
functions.append_todo_functions();
}
@@ -208,6 +218,7 @@ impl Agent {
session_dynamic_instructions: None,
functions,
rag,
graph_rags,
model,
vault: app_state.vault.clone(),
})
@@ -287,10 +298,13 @@ impl Agent {
.display()
.to_string()
.into();
value["config_file"] = paths::agent_config_file(&self.name)
.display()
.to_string()
.into();
let config_path = paths::agent_config_file(&self.name);
let definition_file = if config_path.exists() {
config_path
} else {
paths::agent_graph_file(&self.name)
};
value["config_file"] = definition_file.display().to_string().into();
let data = serde_yaml::to_string(&value)?;
Ok(data)
}
@@ -311,6 +325,10 @@ impl Agent {
self.rag.clone()
}
pub fn graph_rag(&self, node_id: &str) -> Option<Arc<Rag>> {
self.graph_rags.get(node_id).cloned()
}
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
self.functions.append_mcp_meta_functions(mcp_servers);
}
@@ -415,6 +433,14 @@ impl Agent {
self.config.max_auto_continues
}
pub fn inject_todo_instructions(&self) -> bool {
self.config.inject_todo_instructions
}
pub fn continuation_prompt_value(&self) -> Option<String> {
self.config.continuation_prompt.clone()
}
pub fn can_spawn_agents(&self) -> bool {
self.config.can_spawn_agents
}
@@ -439,18 +465,6 @@ impl Agent {
self.config.escalation_timeout
}
pub fn continuation_prompt(&self) -> String {
self.config.continuation_prompt.clone().unwrap_or_else(|| {
formatdoc! {"
[SYSTEM REMINDER - TODO CONTINUATION]
You have incomplete tasks. Rules:
1. BEFORE marking a todo done: verify the work compiles/works. No premature completion.
2. If a todo is broad (e.g. \"implement X and implement Y\"): break it into specific subtasks FIRST using todo__add, then work on those.\n\
3. Each todo should be atomic and be \"single responsibility\" - completable in one focused action.
4. Continue with the next pending item now. Call tools immediately."}
})
}
pub fn compression_threshold(&self) -> Option<usize> {
self.config.compression_threshold
}
@@ -654,6 +668,25 @@ impl AgentConfig {
Ok(agent_config)
}
pub fn from_graph(dir_name: &str, graph: &Graph) -> Self {
AgentConfig {
name: dir_name.to_string(),
model_id: graph.model.clone(),
temperature: graph.temperature,
top_p: graph.top_p,
description: graph.description.clone(),
global_tools: graph.global_tools.clone(),
mcp_servers: graph.mcp_servers.clone(),
conversation_starters: graph.conversation_starters.clone(),
variables: graph.variables.clone(),
can_spawn_agents: graph.has_agent_node(),
max_concurrent_agents: default_max_concurrent_agents(),
max_agent_depth: default_max_agent_depth(),
escalation_timeout: default_escalation_timeout(),
..AgentConfig::default()
}
}
fn load_envs(&mut self, app: &AppConfig) {
let name = &self.name;
let with_prefix = |v: &str| normalize_env_name(&format!("{name}_{v}"));
@@ -750,6 +783,136 @@ pub struct AgentVariable {
pub value: String,
}
fn resolve_document_paths(
documents: &[String],
loaders: &HashMap<String, String>,
agent_data_dir: &Path,
) -> Result<Vec<String>> {
let mut document_paths = vec![];
for path in documents {
if is_url(path) {
document_paths.push(path.to_string());
} else if is_loader_protocol(loaders, path) {
let (protocol, document_path) = path
.split_once(':')
.with_context(|| "Invalid loader protocol path")?;
let resolved_path = resolve_home_dir(document_path);
let new_path = if Path::new(&resolved_path).is_relative() {
safe_join_path(agent_data_dir, resolved_path)
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?
} else {
PathBuf::from(&resolved_path)
};
document_paths.push(format!("{}:{}", protocol, new_path.display()));
} else if Path::new(&resolve_home_dir(path)).is_relative() {
let new_path = safe_join_path(agent_data_dir, path)
.ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?;
document_paths.push(new_path.display().to_string())
} else {
document_paths.push(path.to_string())
}
}
Ok(document_paths)
}
#[allow(clippy::too_many_arguments)]
async fn init_graph_rags(
app: &AppConfig,
app_state: &AppState,
agent_name: &str,
graph: &Graph,
agent_data_dir: &Path,
loaders: &HashMap<String, String>,
info_flag: bool,
abort_signal: AbortSignal,
) -> Result<HashMap<String, Arc<Rag>>> {
let mut rags = HashMap::new();
if info_flag {
return Ok(rags);
}
for (node_id, node) in &graph.nodes {
let NodeType::Rag(rag_node) = &node.node_type else {
continue;
};
let rag_path = paths::agent_rag_file(agent_name, node_id);
let key = RagKey::GraphNode {
agent: agent_name.to_string(),
node: node_id.clone(),
};
let rag = if rag_path.exists() {
let app_clone = app.clone();
let path_clone = rag_path.clone();
let name_clone = node_id.clone();
app_state
.rag_cache
.load_with(key, || async move {
Rag::load(&app_clone, &name_clone, &path_clone)
})
.await?
} else {
let config = RagInitConfig {
embedding_model: rag_node.embedding_model.clone(),
chunk_size: rag_node.chunk_size,
chunk_overlap: rag_node.chunk_overlap,
reranker_model: rag_node.reranker_model.clone(),
top_k: rag_node.top_k,
batch_size: rag_node.batch_size,
};
let fully_specified = config.embedding_model.is_some()
&& config.chunk_size.is_some()
&& config.chunk_overlap.is_some();
if !fully_specified {
if !*IS_STDOUT_TERMINAL {
bail!(
"Agent '{agent_name}' requires RAG for rag node '{node_id}', but its \
knowledge base is not built and the node does not fully specify how \
to build it. Set `embedding_model`, `chunk_size`, and `chunk_overlap` \
on the node, or run the agent once interactively."
);
}
let ans = Confirm::new(&format!(
"Initialize RAG knowledge base for rag node '{node_id}'?"
))
.with_default(true)
.prompt()?;
if !ans {
bail!(
"Agent '{agent_name}' has rag node '{node_id}' but its RAG was not \
initialized. RAG initialization is required for this agent."
);
}
}
let document_paths =
resolve_document_paths(&rag_node.documents, loaders, agent_data_dir)?;
let app_clone = app.clone();
let path_clone = rag_path.clone();
let name_clone = node_id.clone();
let abort = abort_signal.clone();
app_state
.rag_cache
.load_with(key, || async move {
Rag::init_with_config(
&app_clone,
&name_clone,
&path_clone,
&document_paths,
&config,
abort,
)
.await
})
.await?
};
rags.insert(node_id.clone(), rag);
}
Ok(rags)
}
pub fn list_agents() -> Vec<String> {
let agents_data_dir = paths::agents_data_dir();
if !agents_data_dir.exists() {
@@ -876,4 +1039,88 @@ variables:
assert!(config.inject_todo_instructions);
assert!(config.inject_spawn_instructions);
}
#[test]
fn from_graph_maps_agent_level_fields() {
let yaml = formatdoc! {r#"
name: graph_name_ignored
description: A graph agent
model: claude:claude-sonnet-4-6
temperature: 0.3
top_p: 0.8
global_tools:
- fetch_pdf.sh
mcp_servers:
- pubmed-search
conversation_starters:
- "Start here"
start: e
nodes:
e:
id: e
type: end
output: done
"#};
let graph: Graph = serde_yaml::from_str(&yaml).unwrap();
let config = AgentConfig::from_graph("my-agent-dir", &graph);
assert_eq!(config.name, "my-agent-dir");
assert_eq!(config.description, "A graph agent");
assert_eq!(config.model_id.as_deref(), Some("claude:claude-sonnet-4-6"));
assert_eq!(config.temperature, Some(0.3));
assert_eq!(config.top_p, Some(0.8));
assert_eq!(config.global_tools, vec!["fetch_pdf.sh"]);
assert_eq!(config.mcp_servers, vec!["pubmed-search"]);
assert_eq!(config.conversation_starters, vec!["Start here"]);
}
#[test]
fn from_graph_derives_can_spawn_agents_from_agent_nodes() {
let with_agent = formatdoc! {r#"
name: g
start: a
nodes:
a:
id: a
type: agent
agent: helper
prompt: hi
next: e
e:
id: e
type: end
output: done
"#};
let graph: Graph = serde_yaml::from_str(&with_agent).unwrap();
assert!(AgentConfig::from_graph("d", &graph).can_spawn_agents);
let no_agent =
"name: g\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
let graph: Graph = serde_yaml::from_str(no_agent).unwrap();
assert!(!AgentConfig::from_graph("d", &graph).can_spawn_agents);
}
#[test]
fn from_graph_keeps_defaults_for_llm_loop_fields() {
let yaml = "name: g\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
let config = AgentConfig::from_graph("d", &graph);
assert!(!config.auto_continue);
assert!(config.instructions.is_empty());
assert!(config.documents.is_empty());
assert!(!config.inject_todo_instructions);
assert!(!config.inject_spawn_instructions);
assert_eq!(config.max_auto_continues, 0);
assert_eq!(config.summarization_threshold, 0);
assert_eq!(
config.max_concurrent_agents,
default_max_concurrent_agents()
);
assert_eq!(config.max_agent_depth, default_max_agent_depth());
assert_eq!(config.escalation_timeout, default_escalation_timeout());
}
}
+15
View File
@@ -39,6 +39,11 @@ pub struct AppConfig {
pub mapping_mcp_servers: IndexMap<String, String>,
pub enabled_mcp_servers: Option<String>,
pub auto_continue: bool,
pub max_auto_continues: usize,
pub inject_todo_instructions: bool,
pub continuation_prompt: Option<String>,
pub repl_prelude: Option<String>,
pub cmd_prelude: Option<String>,
pub agent_session: Option<String>,
@@ -95,6 +100,11 @@ impl Default for AppConfig {
mapping_mcp_servers: Default::default(),
enabled_mcp_servers: None,
auto_continue: false,
max_auto_continues: 10,
inject_todo_instructions: true,
continuation_prompt: None,
repl_prelude: None,
cmd_prelude: None,
agent_session: None,
@@ -152,6 +162,11 @@ impl AppConfig {
mapping_mcp_servers: config.mapping_mcp_servers,
enabled_mcp_servers: config.enabled_mcp_servers,
auto_continue: config.auto_continue,
max_auto_continues: config.max_auto_continues,
inject_todo_instructions: config.inject_todo_instructions,
continuation_prompt: config.continuation_prompt,
repl_prelude: config.repl_prelude,
cmd_prelude: config.cmd_prelude,
agent_session: config.agent_session,
File diff suppressed because it is too large Load Diff
+2 -2
View File
@@ -85,7 +85,7 @@ impl Macro {
Ok(value)
}
pub fn install_macros() -> Result<()> {
pub fn install_macros(force: bool) -> Result<()> {
info!(
"Installing built-in macros in {}",
paths::macros_dir().display()
@@ -98,7 +98,7 @@ impl Macro {
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
let file_path = paths::macros_dir().join(file.as_ref());
if file_path.exists() {
if file_path.exists() && !force {
debug!(
"Macro file already exists, skipping: {}",
file_path.display()
+7 -6
View File
@@ -109,12 +109,13 @@ impl McpFactory {
mod tests {
use super::*;
use crate::mcp::{JsonField, McpServer, McpTransportType};
use indexmap::IndexMap;
use std::collections::HashMap;
fn stdio_spec(
command: &str,
args: Option<Vec<String>>,
env: Option<HashMap<String, JsonField>>,
env: Option<IndexMap<String, JsonField>>,
) -> McpServer {
McpServer {
transport_type: McpTransportType::Stdio,
@@ -130,7 +131,7 @@ mod tests {
fn remote_spec(
transport: McpTransportType,
url: &str,
headers: Option<HashMap<String, String>>,
headers: Option<IndexMap<String, String>>,
) -> McpServer {
McpServer {
transport_type: transport,
@@ -145,7 +146,7 @@ mod tests {
#[test]
fn key_from_stdio_spec_captures_command_args_env() {
let mut env = HashMap::new();
let mut env = IndexMap::new();
env.insert("TOKEN".into(), JsonField::Str("abc".into()));
let spec = stdio_spec("npx", Some(vec!["-y".into(), "server".into()]), Some(env));
let key = McpServerKey::from_spec("my-server", &spec);
@@ -163,7 +164,7 @@ mod tests {
#[test]
fn key_from_stdio_spec_sorts_args_and_env() {
let mut env = HashMap::new();
let mut env = IndexMap::new();
env.insert("Z_VAR".into(), JsonField::Str("z".into()));
env.insert("A_VAR".into(), JsonField::Int(42));
let spec = stdio_spec(
@@ -222,7 +223,7 @@ mod tests {
#[test]
fn key_from_remote_sse_spec_with_sorted_headers() {
let mut hdrs = HashMap::new();
let mut hdrs = IndexMap::new();
hdrs.insert("Z-Key".into(), "z-val".into());
hdrs.insert("A-Key".into(), "a-val".into());
let spec = remote_spec(McpTransportType::Sse, "http://sse.example.com", Some(hdrs));
@@ -264,7 +265,7 @@ mod tests {
#[test]
fn key_env_bool_and_int_coerce_to_string() {
let mut env = HashMap::new();
let mut env = IndexMap::new();
env.insert("FLAG".into(), JsonField::Bool(true));
env.insert("PORT".into(), JsonField::Int(3000));
let spec = stdio_spec("cmd", None, Some(env));
+152 -5
View File
@@ -2,6 +2,7 @@ mod agent;
mod app_config;
mod app_state;
mod input;
mod install_remote;
mod macros;
mod mcp_factory;
pub(crate) mod paths;
@@ -12,19 +13,24 @@ mod role;
mod session;
pub(crate) mod todo;
mod tool_scope;
mod update;
pub use self::agent::{Agent, AgentVariables, complete_agent_variables, list_agents};
pub use self::agent::{
Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents,
};
#[allow(unused_imports)]
pub use self::app_config::AppConfig;
#[allow(unused_imports)]
pub use self::app_state::AppState;
pub use self::input::Input;
pub use self::install_remote::{install_remote, install_remote_from_repl_args};
#[allow(unused_imports)]
pub use self::request_context::RequestContext;
pub use self::request_context::{RenderMode, RequestContext};
pub use self::role::{
CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, Role, RoleLike, SHELL_ROLE,
};
use self::session::Session;
pub use self::update::run_self_update;
use crate::client::{
ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
ProviderModels, create_client_config, list_client_types,
@@ -66,6 +72,7 @@ const DARK_THEME: &[u8] = include_bytes!("../../assets/monokai-extended.theme.bi
const LIGHT_THEME: &[u8] = include_bytes!("../../assets/monokai-extended-light.theme.bin");
const CONFIG_FILE_NAME: &str = "config.yaml";
const AGENT_GRAPH_FILE_NAME: &str = "graph.yaml";
const ROLES_DIR_NAME: &str = "roles";
const MACROS_DIR_NAME: &str = "macros";
const ENV_FILE_NAME: &str = ".env";
@@ -79,6 +86,26 @@ const GLOBAL_TOOLS_DIR_NAME: &str = "tools";
const GLOBAL_TOOLS_UTILS_DIR_NAME: &str = "utils";
const BASH_PROMPT_UTILS_FILE_NAME: &str = "prompt-utils.sh";
const MCP_FILE_NAME: &str = "mcp.json";
const DEFAULT_VISIBLE_TOOLS: [&str; 18] = [
"execute_command.sh",
"execute_py_code.py",
"execute_sql_code.sh",
"fetch_url_via_curl.sh",
"fs_cat.sh",
"fs_glob.sh",
"fs_grep.sh",
"fs_ls.sh",
"fs_mkdir.sh",
"fs_patch.sh",
"fs_read.sh",
"fs_rm.sh",
"fs_write.sh",
"get_current_time.sh",
"get_current_weather.sh",
"search_wikipedia.sh",
"search_arxiv.sh",
"web_search_loki.sh",
];
const CLIENTS_FIELD: &str = "clients";
@@ -121,6 +148,11 @@ pub struct Config {
pub mapping_mcp_servers: IndexMap<String, String>,
pub enabled_mcp_servers: Option<String>,
pub auto_continue: bool,
pub max_auto_continues: usize,
pub inject_todo_instructions: bool,
pub continuation_prompt: Option<String>,
pub repl_prelude: Option<String>,
pub cmd_prelude: Option<String>,
pub agent_session: Option<String>,
@@ -177,6 +209,11 @@ impl Default for Config {
mapping_mcp_servers: Default::default(),
enabled_mcp_servers: None,
auto_continue: false,
max_auto_continues: 10,
inject_todo_instructions: true,
continuation_prompt: None,
repl_prelude: None,
cmd_prelude: None,
agent_session: None,
@@ -210,12 +247,110 @@ impl Default for Config {
}
pub fn install_builtins() -> Result<()> {
Functions::install_builtin_global_tools()?;
Agent::install_builtin_agents()?;
Macro::install_macros()?;
Functions::install_builtin_global_tools(false)?;
Agent::install_builtin_agents(false)?;
Macro::install_macros(false)?;
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum AssetCategory {
Agents,
Macros,
Functions,
#[value(name = "mcp_config")]
McpConfig,
}
impl AssetCategory {
pub const NAMES: [&'static str; 4] = ["agents", "macros", "functions", "mcp_config"];
pub fn parse(name: &str) -> Option<Self> {
match name {
"agents" => Some(Self::Agents),
"macros" => Some(Self::Macros),
"functions" => Some(Self::Functions),
"mcp_config" => Some(Self::McpConfig),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
pub enum InstallFilter {
Agents,
Roles,
Macros,
Functions,
#[value(name = "mcp_config")]
McpConfig,
}
impl InstallFilter {
pub const NAMES: [&'static str; 5] = ["agents", "roles", "macros", "functions", "mcp_config"];
pub fn parse(name: &str) -> Option<Self> {
match name {
"agents" => Some(Self::Agents),
"roles" => Some(Self::Roles),
"macros" => Some(Self::Macros),
"functions" => Some(Self::Functions),
"mcp_config" => Some(Self::McpConfig),
_ => None,
}
}
}
pub fn install_assets(category: AssetCategory) -> Result<()> {
let (label, target) = match category {
AssetCategory::Agents => ("agents", paths::agents_data_dir()),
AssetCategory::Macros => ("macros", paths::macros_dir()),
AssetCategory::Functions => ("functions", paths::functions_dir()),
AssetCategory::McpConfig => ("MCP config", paths::mcp_config_file()),
};
if !confirm_asset_overwrite(category, label, &target)? {
println!("Aborted. No files were changed.");
return Ok(());
}
match category {
AssetCategory::Agents => Agent::install_builtin_agents(true)?,
AssetCategory::Macros => Macro::install_macros(true)?,
AssetCategory::Functions => Functions::install_builtin_global_tools(true)?,
AssetCategory::McpConfig => Functions::install_mcp_config()?,
}
println!("Reinstalled bundled {label} ({})", target.display());
Ok(())
}
fn confirm_asset_overwrite(category: AssetCategory, label: &str, target: &Path) -> Result<bool> {
if !*IS_STDOUT_TERMINAL {
return Ok(true);
}
let body = match category {
AssetCategory::McpConfig => format!(
"This replaces your MCP server configuration at {} with this \
build's bundled template. Your configured MCP servers (and any \
custom secret references they contain) will be lost.",
target.display()
),
_ => format!(
"Reinstalling bundled {label} overwrites every bundled {label} in \
{} with this build's packaged versions. Local changes to bundled \
{label} will be lost; {label} you created yourself are left \
untouched.",
target.display()
),
};
let prompt = format!("{} {body}\nContinue? [y/N] ", warning_text("WARNING:"));
let answer = read_single_key(&['y', 'Y', 'n', 'N'], 'n', &prompt)?;
Ok(matches!(answer, 'y' | 'Y'))
}
pub fn default_sessions_dir() -> PathBuf {
match env::var(get_env_name("sessions_dir")) {
Ok(value) => PathBuf::from(value),
@@ -474,6 +609,18 @@ pub async fn create_config_file(config_path: &Path) -> Result<()> {
let (model, clients_config) = create_client_config(client, &vault).await?;
config["model"] = model.into();
config["vault_password_file"] = vault.password_file()?.display().to_string().into();
config["stream"] = json!(true);
config["save"] = json!(true);
config["keybindings"] = json!("vi");
config["wrap"] = json!("auto");
config["wrap_code"] = json!(false);
config["function_calling_support"] = json!(true);
config["enabled_tools"] = json!(null);
config["visible_tools"] = json!(DEFAULT_VISIBLE_TOOLS);
config["mcp_server_support"] = json!(true);
config["enabled_mcp_servers"] = json!(null);
config["highlight"] = json!(true);
config["light_theme"] = json!(false);
config[CLIENTS_FIELD] = clients_config;
let config_data = serde_yaml::to_string(&config).with_context(|| "Failed to create config")?;
+8 -3
View File
@@ -1,8 +1,9 @@
use super::role::Role;
use super::{
AGENTS_DIR_NAME, BASH_PROMPT_UTILS_FILE_NAME, CONFIG_FILE_NAME, ENV_FILE_NAME,
FUNCTIONS_BIN_DIR_NAME, FUNCTIONS_DIR_NAME, GLOBAL_TOOLS_DIR_NAME, GLOBAL_TOOLS_UTILS_DIR_NAME,
MACROS_DIR_NAME, MCP_FILE_NAME, ModelsOverride, RAGS_DIR_NAME, ROLES_DIR_NAME,
AGENT_GRAPH_FILE_NAME, AGENTS_DIR_NAME, BASH_PROMPT_UTILS_FILE_NAME, CONFIG_FILE_NAME,
ENV_FILE_NAME, FUNCTIONS_BIN_DIR_NAME, FUNCTIONS_DIR_NAME, GLOBAL_TOOLS_DIR_NAME,
GLOBAL_TOOLS_UTILS_DIR_NAME, MACROS_DIR_NAME, MCP_FILE_NAME, ModelsOverride, RAGS_DIR_NAME,
ROLES_DIR_NAME,
};
use crate::client::ProviderModels;
use crate::utils::{get_env_name, list_file_names, normalize_env_name};
@@ -127,6 +128,10 @@ pub fn agent_data_dir(name: &str) -> PathBuf {
}
}
pub fn agent_graph_file(agent_name: &str) -> PathBuf {
agent_data_dir(agent_name).join(AGENT_GRAPH_FILE_NAME)
}
pub fn agent_config_file(name: &str) -> PathBuf {
match env::var(format!("{}_CONFIG_FILE", normalize_env_name(name))) {
Ok(value) => PathBuf::from(value),
+1
View File
@@ -9,6 +9,7 @@ use std::sync::{Arc, Weak};
pub enum RagKey {
Named(String),
Agent(String),
GraphNode { agent: String, node: String },
}
#[derive(Default)]
+606 -29
View File
@@ -1,14 +1,14 @@
use super::MessageContentToolCalls;
use super::rag_cache::{RagCache, RagKey};
use super::session::Session;
use super::todo::TodoList;
use super::tool_scope::{McpRuntime, ToolScope};
use super::{
AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, CREATE_TITLE_ROLE, Input,
LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, RIGHT_PROMPT, Role, RoleLike, SESSIONS_DIR_NAME,
SUMMARIZATION_PROMPT, SUMMARY_CONTEXT_PROMPT, StateFlags, TEMP_ROLE_NAME, TEMP_SESSION_NAME,
WorkingMode, ensure_parent_exists, list_agents, paths,
AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE,
Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, RIGHT_PROMPT, Role,
RoleLike, SESSIONS_DIR_NAME, SUMMARIZATION_PROMPT, SUMMARY_CONTEXT_PROMPT, StateFlags,
TEMP_ROLE_NAME, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, list_agents, paths,
};
use super::{MessageContentToolCalls, prompts};
use crate::client::{Model, ModelType, list_models};
use crate::function::{
FunctionDeclaration, Functions, ToolCallTracker, ToolResult,
@@ -27,16 +27,33 @@ use crate::utils::{
list_file_names, now, render_prompt, temp_file,
};
use crate::graph;
use anyhow::{Context, Error, Result, bail};
#[cfg(test)]
use indexmap::IndexMap;
use indoc::formatdoc;
use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation};
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use std::env;
use std::fs::{File, OpenOptions, read_dir, read_to_string, remove_dir_all, remove_file};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::{env, fs};
pub struct AutoContinueConfig {
pub enabled: bool,
pub max_continues: usize,
pub inject_instructions: bool,
pub continuation_prompt: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RenderMode {
#[default]
Streaming,
Silent,
}
pub struct RequestContext {
pub app: Arc<AppState>,
@@ -66,6 +83,8 @@ pub struct RequestContext {
pub auto_continue_count: usize,
pub todo_list: TodoList,
pub last_continuation_response: Option<String>,
pub render_mode: RenderMode,
}
impl RequestContext {
@@ -92,6 +111,7 @@ impl RequestContext {
auto_continue_count: 0,
todo_list: TodoList::default(),
last_continuation_response: None,
render_mode: RenderMode::default(),
}
}
@@ -138,9 +158,51 @@ impl RequestContext {
auto_continue_count: 0,
todo_list: TodoList::default(),
last_continuation_response: None,
render_mode: RenderMode::default(),
})
}
/// Forks the context for one parallel branch of a graph super-step.
///
/// Each branch gets a fresh, owned clone. Mutations (role swap,
/// `before/after_chat_completion`, tool tracker, last_message, etc.) are
/// scoped to the branch and discarded when the branch finishes. The
/// user-visible state communication happens through the graph's
/// `StateManager` (via `fork_for_branch_state` + `diff_against` +
/// `apply_branch_writes` reducers), and not through `RequestContext`.
///
/// Distinction from `new_for_child`: `new_for_child` builds a fresh context
/// for a spawned sub-agent (different agent identity, different supervisor
/// hierarchy, depth+1, fresh tool tracker). `fork_for_branch` keeps the
/// caller's identity and supervisor hierarchy; it's a sibling clone of the
/// same logical agent, running one of N parallel work items.
pub fn fork_for_branch(&self) -> Self {
Self {
app: Arc::clone(&self.app),
macro_flag: self.macro_flag,
info_flag: self.info_flag,
working_mode: self.working_mode,
model: self.model.clone(),
agent_variables: self.agent_variables.clone(),
role: self.role.clone(),
session: self.session.clone(),
rag: self.rag.clone(),
agent: self.agent.clone(),
last_message: self.last_message.clone(),
tool_scope: self.tool_scope.clone(),
supervisor: self.supervisor.clone(),
parent_supervisor: self.parent_supervisor.clone(),
self_agent_id: self.self_agent_id.clone(),
inbox: self.inbox.clone(),
escalation_queue: self.escalation_queue.clone(),
current_depth: self.current_depth,
auto_continue_count: 0,
todo_list: self.todo_list.clone(),
last_continuation_response: None,
render_mode: self.render_mode,
}
}
pub fn new_for_child(
app: Arc<AppState>,
parent: &Self,
@@ -176,6 +238,7 @@ impl RequestContext {
auto_continue_count: 0,
todo_list: TodoList::default(),
last_continuation_response: None,
render_mode: parent.render_mode,
}
}
@@ -523,7 +586,7 @@ impl RequestContext {
}
pub fn extract_role(&self, app: &AppConfig) -> Role {
if let Some(session) = self.session.as_ref() {
let mut role = if let Some(session) = self.session.as_ref() {
session.to_role()
} else if let Some(agent) = self.agent.as_ref() {
agent.to_role()
@@ -539,6 +602,65 @@ impl RequestContext {
app.enabled_mcp_servers.clone(),
);
role
};
if self.agent.is_none() && self.app.config.function_calling_support {
let config = self.auto_continue_config();
if config.enabled && config.inject_instructions {
role.append_to_prompt(prompts::DEFAULT_TODO_INSTRUCTIONS);
}
}
role
}
pub fn auto_continue_config(&self) -> AutoContinueConfig {
if let Some(agent) = &self.agent {
return AutoContinueConfig {
enabled: agent.auto_continue_enabled(),
max_continues: agent.max_auto_continues(),
inject_instructions: agent.inject_todo_instructions(),
continuation_prompt: agent.continuation_prompt_value(),
};
}
let app = &self.app.config;
let enabled = self
.session
.as_ref()
.and_then(|s| s.auto_continue())
.or_else(|| self.role.as_ref().and_then(|r| r.auto_continue()))
.unwrap_or(app.auto_continue);
let max = self
.session
.as_ref()
.and_then(|s| s.max_auto_continues())
.or_else(|| self.role.as_ref().and_then(|r| r.max_auto_continues()))
.unwrap_or(app.max_auto_continues);
let inject = self
.session
.as_ref()
.and_then(|s| s.inject_todo_instructions())
.or_else(|| {
self.role
.as_ref()
.and_then(|r| r.inject_todo_instructions())
})
.unwrap_or(app.inject_todo_instructions);
let prompt = self
.session
.as_ref()
.and_then(|s| s.continuation_prompt().map(|v| v.to_string()))
.or_else(|| {
self.role
.as_ref()
.and_then(|r| r.continuation_prompt().map(|v| v.to_string()))
})
.or_else(|| app.continuation_prompt.clone());
AutoContinueConfig {
enabled,
max_continues: max,
inject_instructions: inject,
continuation_prompt: prompt,
}
}
@@ -747,6 +869,8 @@ impl RequestContext {
app.function_calling_support.to_string(),
),
("mcp_server_support", app.mcp_server_support.to_string()),
("auto_continue", app.auto_continue.to_string()),
("max_auto_continues", app.max_auto_continues.to_string()),
("stream", app.stream.to_string()),
("save", app.save.to_string()),
("keybindings", app.keybindings.clone()),
@@ -922,9 +1046,12 @@ impl RequestContext {
let app = self.app.config.as_ref();
let mut functions = vec![];
if app.function_calling_support {
if let Some(enabled_tools) = role.enabled_tools() {
let mut tool_names: HashSet<String> = Default::default();
let declaration_names: HashSet<String> = self
// Compute the set of tool names enabled by the role filter, drawn
// from BOTH the tool_scope pool and the agent's pool so that an
// explicit `enabled_tools` list (e.g. from a graph LLM node) can
// narrow the agent's own custom tools too.
let role_filter: Option<HashSet<String>> = role.enabled_tools().map(|enabled_tools| {
let mut declaration_names: HashSet<String> = self
.tool_scope
.functions
.declarations()
@@ -936,11 +1063,32 @@ impl RequestContext {
})
.map(|v| v.name.to_string())
.collect();
if let Some(agent) = &self.agent {
declaration_names.extend(
agent
.functions()
.declarations()
.iter()
.filter(|v| {
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
&& !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
})
.map(|v| v.name.to_string()),
);
}
let mut tool_names: HashSet<String> = Default::default();
if enabled_tools == "all" {
tool_names.extend(declaration_names);
} else {
for item in enabled_tools.split(',') {
let item = item.trim();
if item.is_empty() {
continue;
}
if let Some(values) = app.mapping_tools.get(item) {
tool_names.extend(
values
@@ -953,6 +1101,10 @@ impl RequestContext {
}
}
}
tool_names
});
if let Some(ref tool_names) = role_filter {
functions = self
.tool_scope
.functions
@@ -995,6 +1147,11 @@ impl RequestContext {
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
})
.collect();
if let Some(ref tool_names) = role_filter {
agent_functions.retain(|v| tool_names.contains(&v.name));
}
let tool_names: HashSet<String> = agent_functions
.iter()
.filter_map(|v| {
@@ -1021,9 +1178,9 @@ impl RequestContext {
let app = self.app.config.as_ref();
let mut mcp_functions = vec![];
if app.mcp_server_support {
if let Some(enabled_mcp_servers) = role.enabled_mcp_servers() {
let mut server_names: HashSet<String> = Default::default();
let mcp_declaration_names: HashSet<String> = self
let role_filter: Option<HashSet<String>> =
role.enabled_mcp_servers().map(|enabled_mcp_servers| {
let mut mcp_declaration_names: HashSet<String> = self
.tool_scope
.functions
.declarations()
@@ -1035,11 +1192,32 @@ impl RequestContext {
})
.map(|v| v.name.to_string())
.collect();
if let Some(agent) = &self.agent {
mcp_declaration_names.extend(
agent
.functions()
.declarations()
.iter()
.filter(|v| {
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|| v.name
.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
})
.map(|v| v.name.to_string()),
);
}
let mut server_names: HashSet<String> = Default::default();
if enabled_mcp_servers == "all" {
server_names.extend(mcp_declaration_names);
} else {
for item in enabled_mcp_servers.split(',') {
let item = item.trim();
if item.is_empty() {
continue;
}
let item_invoke_name =
format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
let item_search_name =
@@ -1078,6 +1256,10 @@ impl RequestContext {
}
}
}
server_names
});
if let Some(ref server_names) = role_filter {
mcp_functions = self
.tool_scope
.functions
@@ -1105,6 +1287,11 @@ impl RequestContext {
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
})
.collect();
if let Some(ref server_names) = role_filter {
agent_functions.retain(|v| server_names.contains(&v.name));
}
let tool_names: HashSet<String> = agent_functions
.iter()
.filter_map(|v| {
@@ -1212,6 +1399,19 @@ impl RequestContext {
Ok(())
}
pub fn edit_mcp_config(&self) -> Result<()> {
let mcp_path = paths::mcp_config_file();
let editor = self.app.config.editor()?;
edit_file(&editor, &mcp_path)?;
println!(
"NOTE: Remember to restart {} for changes to '{}' to take effect",
env!("CARGO_CRATE_NAME"),
mcp_path.display(),
);
Ok(())
}
pub fn new_role(&self, app: &AppConfig, name: &str) -> Result<()> {
if self.macro_flag {
bail!("No role");
@@ -1286,21 +1486,30 @@ impl RequestContext {
Some(agent) => agent.name(),
None => bail!("No agent"),
};
let agent_config_path = paths::agent_config_file(agent_name);
ensure_parent_exists(&agent_config_path)?;
if !agent_config_path.exists() {
std::fs::write(
&agent_config_path,
let config_path = paths::agent_config_file(agent_name);
let graph_path = paths::agent_graph_file(agent_name);
let target_path = if !config_path.exists() && graph_path.exists() {
graph_path
} else {
config_path
};
ensure_parent_exists(&target_path)?;
if !target_path.exists() {
fs::write(
&target_path,
"# see https://github.com/Dark-Alex-17/loki/blob/main/config.agent.example.yaml\n",
)
.with_context(|| format!("Failed to write to '{}'", agent_config_path.display()))?;
.with_context(|| format!("Failed to write to '{}'", target_path.display()))?;
}
let editor = app.editor()?;
edit_file(&editor, &agent_config_path)?;
edit_file(&editor, &target_path)?;
println!(
"NOTE: Remember to reload the agent if there are changes made to '{}'",
agent_config_path.display()
target_path.display()
);
Ok(())
}
@@ -1402,12 +1611,24 @@ impl RequestContext {
}
pub async fn update(&mut self, data: &str, abort_signal: AbortSignal) -> Result<()> {
let parts: Vec<&str> = data.split_whitespace().collect();
if parts.len() != 2 {
let (key, raw_value) = match data.split_once(char::is_whitespace) {
Some((k, v)) => (k, v.trim()),
None => bail!("Usage: .set <key> <value>. If value is null, unset key."),
};
if raw_value.is_empty() {
bail!("Usage: .set <key> <value>. If value is null, unset key.");
}
let key = parts[0];
let value = parts[1];
let value = match key {
"continuation_prompt" => raw_value,
_ => {
if raw_value.contains(char::is_whitespace) {
bail!("Usage: .set <key> <value>. If value is null, unset key.");
}
raw_value
}
};
match key {
"temperature" => {
let value = super::parse_value(value)?;
@@ -1522,6 +1743,49 @@ impl RequestContext {
let value = value.parse().with_context(|| "Invalid value")?;
self.update_app_config(|app| app.highlight = value);
}
"auto_continue" => {
let value: bool = value.parse().with_context(|| "Invalid value")?;
if value && !self.app.config.function_calling_support {
bail!(
"Cannot enable auto_continue: function calling is disabled. Set 'function_calling_support: true' first."
);
}
if let Some(session) = self.session.as_mut() {
session.set_auto_continue(Some(value));
} else {
self.update_app_config(|app| app.auto_continue = value);
}
if value
&& self.app.config.function_calling_support
&& !self.tool_scope.functions.contains("todo__init")
{
self.tool_scope.functions.append_todo_functions();
}
}
"max_auto_continues" => {
let value: usize = value.parse().with_context(|| "Invalid value")?;
if let Some(session) = self.session.as_mut() {
session.set_max_auto_continues(Some(value));
} else {
self.update_app_config(|app| app.max_auto_continues = value);
}
}
"inject_todo_instructions" => {
let value: bool = value.parse().with_context(|| "Invalid value")?;
if let Some(session) = self.session.as_mut() {
session.set_inject_todo_instructions(Some(value));
} else {
self.update_app_config(|app| app.inject_todo_instructions = value);
}
}
"continuation_prompt" => {
let value: Option<String> = super::parse_value(value)?;
if let Some(session) = self.session.as_mut() {
session.set_continuation_prompt(value);
} else {
self.update_app_config(|app| app.continuation_prompt = value);
}
}
_ => bail!("Unknown key '{key}'"),
}
Ok(())
@@ -1591,6 +1855,12 @@ impl RequestContext {
}
".rag" => super::map_completion_values(paths::list_rags()),
".agent" => super::map_completion_values(list_agents()),
".install" => {
let mut values: Vec<String> =
AssetCategory::NAMES.iter().map(|s| s.to_string()).collect();
values.push("remote".to_string());
super::map_completion_values(values)
}
".macro" => super::map_completion_values(paths::list_macros()),
".starter" => match &self.agent {
Some(agent) => agent
@@ -1603,10 +1873,14 @@ impl RequestContext {
},
".set" => {
let mut values = vec![
"auto_continue",
"continuation_prompt",
"temperature",
"top_p",
"enabled_tools",
"enabled_mcp_servers",
"inject_todo_instructions",
"max_auto_continues",
"save_session",
"compression_threshold",
"rag_reranker_model",
@@ -1642,6 +1916,28 @@ impl RequestContext {
}
_ => vec![],
};
} else if cmd == ".install" && args.first() == Some(&"remote") && args.len() >= 2 {
let prev = args.get(args.len() - 2).copied().unwrap_or("");
if prev == "--filter" {
values = super::map_completion_values(
InstallFilter::NAMES.iter().map(|s| s.to_string()).collect(),
);
} else {
let has_filter = args.iter().enumerate().any(|(i, a)| {
a.starts_with("--filter=") || (*a == "--filter" && i < args.len() - 1)
});
let has_force = args.contains(&"--force");
let mut available: Vec<&str> = vec![];
if !has_filter {
available.push("--filter");
}
if !has_force {
available.push("--force");
}
values = super::map_completion_values(available);
}
} else if cmd == ".set" && args.len() == 2 {
let candidates = match args[0] {
"max_output_tokens" => match self.current_model().max_output_tokens() {
@@ -1721,6 +2017,19 @@ impl RequestContext {
.map(|v| v.id())
.collect(),
"highlight" => super::complete_bool(app.highlight),
"auto_continue" => {
let config = self.auto_continue_config();
super::complete_bool(config.enabled)
}
"max_auto_continues" => {
let config = self.auto_continue_config();
vec![config.max_continues.to_string()]
}
"inject_todo_instructions" => {
let config = self.auto_continue_config();
super::complete_bool(config.inject_instructions)
}
"continuation_prompt" => vec!["null".to_string()],
_ => vec![],
};
values = candidates.into_iter().map(|v| (v, None)).collect();
@@ -1735,7 +2044,7 @@ impl RequestContext {
.collect();
} else if cmd == ".agent" {
if args.len() == 2 {
let dir = paths::agent_data_dir(args[0]).join(super::SESSIONS_DIR_NAME);
let dir = paths::agent_data_dir(args[0]).join(SESSIONS_DIR_NAME);
values = list_file_names(dir, ".yaml")
.into_iter()
.map(|v| (v, None))
@@ -1810,6 +2119,12 @@ impl RequestContext {
if self.working_mode.is_repl() {
functions.append_user_interaction_functions();
}
if self.agent.is_none()
&& app.function_calling_support
&& self.auto_continue_config().enabled
{
functions.append_todo_functions();
}
if !mcp_runtime.is_empty() {
functions.append_mcp_meta_functions(mcp_runtime.server_names());
}
@@ -1830,6 +2145,10 @@ impl RequestContext {
abort_signal: AbortSignal,
) -> Result<()> {
let role = self.retrieve_role(app, name)?;
if let Some(session) = self.session.as_mut() {
session.guard_empty()?;
}
let mcp_servers = if app.mcp_server_support {
role.enabled_mcp_servers()
} else {
@@ -1956,6 +2275,14 @@ impl RequestContext {
)
.await?;
let is_graph_agent = graph::agent_has_graph(agent_name);
if is_graph_agent && session_name.is_some() {
bail!(
"Graph-based agent '{agent_name}' does not support sessions. \
The graph manages its own state; re-run without a session."
);
}
let mcp_servers = if app.mcp_server_support {
(!agent.mcp_server_names().is_empty()).then(|| agent.mcp_server_names().join(","))
} else {
@@ -1977,14 +2304,22 @@ impl RequestContext {
);
}
// Graph agents manage their own state; never engage a session,
// not even an inherited app-level `agent_session` default.
let session_name = session_name.map(|v| v.to_string()).or_else(|| {
if self.macro_flag {
if self.macro_flag || is_graph_agent {
None
} else {
agent.agent_session().map(|v| v.to_string())
}
});
if self.session.is_some() {
bail!(
"Already in a session, please run '.exit session' first to exit the current session."
);
}
let should_init_supervisor = agent.can_spawn_agents();
let max_concurrent = agent.max_concurrent_agents();
let max_depth = agent.max_agent_depth();
@@ -2196,7 +2531,7 @@ impl RequestContext {
.clone()
.unwrap_or_else(|| SUMMARY_CONTEXT_PROMPT.into());
let todo_prefix = if self.agent.is_some() && !self.todo_list.is_empty() {
let todo_prefix = if self.auto_continue_config().enabled && !self.todo_list.is_empty() {
format!(
"[ACTIVE TODO LIST]\n{}\n\n",
self.todo_list.render_for_model()
@@ -2591,7 +2926,7 @@ mod tests {
let mcp_config = if server_names.is_empty() {
None
} else {
let mut servers = HashMap::new();
let mut servers = IndexMap::new();
for name in server_names {
servers.insert(
name.to_string(),
@@ -3289,4 +3624,246 @@ mod tests {
create_dir_all(&rags_dir).unwrap();
assert!(paths::list_rags().is_empty());
}
#[test]
#[serial]
fn use_agent_errors_when_already_in_session() {
let _guard = TestConfigDirGuard::new();
let mut ctx = create_test_ctx();
ctx.session = Some(Session::default());
let app = ctx.app.config.clone();
let agent_name = format!(
"test_agent_{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let agent_dir = paths::agent_data_dir(&agent_name);
create_dir_all(&agent_dir).unwrap();
write(
agent_dir.join("config.yaml"),
format!("name: {agent_name}\ninstructions: hi\n"),
)
.unwrap();
let abort = utils::create_abort_signal();
let result = run_async(ctx.use_agent(&app, &agent_name, Some("test_session"), abort));
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("Already in a session")
);
assert!(
ctx.agent.is_none(),
"Agent should not be set when session check fails"
);
}
#[test]
#[serial]
fn use_agent_errors_when_already_in_session_even_without_session_name() {
let _guard = TestConfigDirGuard::new();
let mut ctx = create_test_ctx();
ctx.session = Some(Session::default());
let app = ctx.app.config.clone();
let agent_name = format!(
"test_agent_{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let agent_dir = paths::agent_data_dir(&agent_name);
create_dir_all(&agent_dir).unwrap();
write(
agent_dir.join("config.yaml"),
format!("name: {agent_name}\ninstructions: hi\n"),
)
.unwrap();
let abort = utils::create_abort_signal();
let result = run_async(ctx.use_agent(&app, &agent_name, None, abort));
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("Already in a session")
);
assert!(
ctx.agent.is_none(),
"Agent should not be set when session check fails"
);
}
#[test]
#[serial]
fn use_agent_errors_when_graph_agent_given_explicit_session() {
let _guard = TestConfigDirGuard::new();
let mut ctx = create_test_ctx();
let app = ctx.app.config.clone();
let agent_name = format!(
"test_graph_agent_{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let agent_dir = paths::agent_data_dir(&agent_name);
create_dir_all(&agent_dir).unwrap();
write(
agent_dir.join("graph.yaml"),
format!(
"name: {agent_name}\nversion: \"1.0\"\nstart: done\nnodes:\n done:\n type: end\n output: ok\n"
),
)
.unwrap();
let abort = utils::create_abort_signal();
let result = run_async(ctx.use_agent(&app, &agent_name, Some("test_session"), abort));
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("does not support sessions")
);
assert!(
ctx.agent.is_none(),
"Agent should not be set when the graph-agent session guard fails"
);
}
#[test]
#[serial]
fn use_agent_skips_inherited_session_for_graph_agent() {
let _guard = TestConfigDirGuard::new();
let mut ctx = create_test_ctx();
ctx.update_app_config(|app| app.agent_session = Some("inherited".to_string()));
let app = ctx.app.config.clone();
let agent_name = format!(
"test_graph_agent_{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let agent_dir = paths::agent_data_dir(&agent_name);
create_dir_all(&agent_dir).unwrap();
write(
agent_dir.join("graph.yaml"),
format!(
"name: {agent_name}\nversion: \"1.0\"\nstart: done\nnodes:\n done:\n type: end\n output: ok\n"
),
)
.unwrap();
let abort = utils::create_abort_signal();
run_async(ctx.use_agent(&app, &agent_name, None, abort)).unwrap();
assert!(ctx.agent.is_some(), "Graph agent should load successfully");
assert!(
ctx.session.is_none(),
"Graph agent must not engage a session, not even an inherited default"
);
}
fn first_file(dir: &Path) -> Option<PathBuf> {
for entry in read_dir(dir).ok()?.flatten() {
let path = entry.path();
if path.is_dir() {
if let Some(found) = first_file(&path) {
return Some(found);
}
} else {
return Some(path);
}
}
None
}
#[test]
fn asset_category_parse_maps_known_names() {
assert_eq!(AssetCategory::parse("agents"), Some(AssetCategory::Agents));
assert_eq!(AssetCategory::parse("macros"), Some(AssetCategory::Macros));
assert_eq!(
AssetCategory::parse("functions"),
Some(AssetCategory::Functions)
);
assert_eq!(
AssetCategory::parse("mcp_config"),
Some(AssetCategory::McpConfig)
);
assert_eq!(AssetCategory::parse("roles"), None);
assert_eq!(AssetCategory::parse(""), None);
}
#[test]
#[serial]
fn install_builtin_agents_force_overwrites_only_with_force() {
let _guard = TestConfigDirGuard::new();
Agent::install_builtin_agents(false).unwrap();
let file =
first_file(&paths::agents_data_dir()).expect("bundled agents should be installed");
write(&file, "SENTINEL").unwrap();
Agent::install_builtin_agents(false).unwrap();
assert_eq!(
read_to_string(&file).unwrap(),
"SENTINEL",
"non-force install must not overwrite an existing file"
);
Agent::install_builtin_agents(true).unwrap();
assert_ne!(
read_to_string(&file).unwrap(),
"SENTINEL",
"force install must overwrite the existing file"
);
}
#[test]
#[serial]
fn install_functions_force_preserves_user_mcp_json() {
let _guard = TestConfigDirGuard::new();
Functions::install_builtin_global_tools(false).unwrap();
let mcp = paths::mcp_config_file();
assert!(mcp.exists(), "mcp.json should be installed on first run");
write(&mcp, "USER_MCP_CONFIG").unwrap();
Functions::install_builtin_global_tools(true).unwrap();
assert_eq!(
read_to_string(&mcp).unwrap(),
"USER_MCP_CONFIG",
"force install must NOT overwrite the user's mcp.json"
);
}
#[test]
#[serial]
fn install_mcp_config_overwrites_existing() {
let _guard = TestConfigDirGuard::new();
Functions::install_mcp_config().unwrap();
let mcp = paths::mcp_config_file();
assert!(mcp.exists(), "install_mcp_config should create mcp.json");
write(&mcp, "USER_MCP_CONFIG").unwrap();
Functions::install_mcp_config().unwrap();
assert_ne!(
read_to_string(&mcp).unwrap(),
"USER_MCP_CONFIG",
"install_mcp_config must overwrite the existing mcp.json"
);
}
}
+50
View File
@@ -55,6 +55,14 @@ pub struct Role {
enabled_tools: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
enabled_mcp_servers: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
auto_continue: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
max_auto_continues: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
inject_todo_instructions: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
continuation_prompt: Option<String>,
#[serde(skip)]
model: Model,
@@ -90,6 +98,14 @@ impl Role {
"enabled_mcp_servers" => {
role.enabled_mcp_servers = value.as_str().map(|v| v.to_string())
}
"auto_continue" => role.auto_continue = value.as_bool(),
"max_auto_continues" => {
role.max_auto_continues = value.as_u64().map(|v| v as usize)
}
"inject_todo_instructions" => role.inject_todo_instructions = value.as_bool(),
"continuation_prompt" => {
role.continuation_prompt = value.as_str().map(|v| v.to_string())
}
_ => (),
}
}
@@ -131,6 +147,20 @@ impl Role {
if let Some(enabled_mcp_servers) = self.enabled_mcp_servers() {
metadata.push(format!("enabled_mcp_servers: {enabled_mcp_servers}"));
}
if let Some(auto_continue) = self.auto_continue {
metadata.push(format!("auto_continue: {auto_continue}"));
}
if let Some(max_auto_continues) = self.max_auto_continues {
metadata.push(format!("max_auto_continues: {max_auto_continues}"));
}
if let Some(inject_todo_instructions) = self.inject_todo_instructions {
metadata.push(format!(
"inject_todo_instructions: {inject_todo_instructions}"
));
}
if let Some(continuation_prompt) = &self.continuation_prompt {
metadata.push(format!("continuation_prompt: {continuation_prompt}"));
}
if metadata.is_empty() {
format!("{}\n", self.prompt)
} else if self.prompt.is_empty() {
@@ -225,6 +255,26 @@ impl Role {
self.prompt.contains(INPUT_PLACEHOLDER)
}
pub fn auto_continue(&self) -> Option<bool> {
self.auto_continue
}
pub fn max_auto_continues(&self) -> Option<usize> {
self.max_auto_continues
}
pub fn inject_todo_instructions(&self) -> Option<bool> {
self.inject_todo_instructions
}
pub fn continuation_prompt(&self) -> Option<&str> {
self.continuation_prompt.as_deref()
}
pub fn append_to_prompt(&mut self, text: &str) {
self.prompt.push_str(text);
}
pub fn echo_messages(&self, input: &Input) -> String {
let input_markdown = input.render();
if self.is_empty_prompt() {
+80 -9
View File
@@ -32,6 +32,14 @@ pub struct Session {
save_session: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
compression_threshold: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
auto_continue: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
max_auto_continues: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
inject_todo_instructions: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
continuation_prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
role_name: Option<String>,
@@ -170,6 +178,18 @@ impl Session {
if let Some(save_session) = self.save_session() {
data["save_session"] = save_session.into();
}
if let Some(auto_continue) = self.auto_continue() {
data["auto_continue"] = auto_continue.into();
}
if let Some(max_auto_continues) = self.max_auto_continues() {
data["max_auto_continues"] = max_auto_continues.into();
}
if let Some(inject_todo_instructions) = self.inject_todo_instructions() {
data["inject_todo_instructions"] = inject_todo_instructions.into();
}
if let Some(continuation_prompt) = self.continuation_prompt() {
data["continuation_prompt"] = continuation_prompt.into();
}
let (tokens, percent) = self.tokens_usage();
data["total_tokens"] = tokens.into();
if let Some(max_input_tokens) = self.model().max_input_tokens() {
@@ -225,6 +245,22 @@ impl Session {
items.push(("compression_threshold", compression_threshold.to_string()));
}
if let Some(auto_continue) = self.auto_continue() {
items.push(("auto_continue", auto_continue.to_string()));
}
if let Some(max_auto_continues) = self.max_auto_continues() {
items.push(("max_auto_continues", max_auto_continues.to_string()));
}
if let Some(inject_todo_instructions) = self.inject_todo_instructions() {
items.push((
"inject_todo_instructions",
inject_todo_instructions.to_string(),
));
}
if let Some(continuation_prompt) = self.continuation_prompt() {
items.push(("continuation_prompt", continuation_prompt.to_string()));
}
if let Some(max_input_tokens) = self.model().max_input_tokens() {
items.push(("max_input_tokens", max_input_tokens.to_string()));
}
@@ -335,6 +371,50 @@ impl Session {
}
}
pub fn auto_continue(&self) -> Option<bool> {
self.auto_continue
}
pub fn max_auto_continues(&self) -> Option<usize> {
self.max_auto_continues
}
pub fn set_auto_continue(&mut self, value: Option<bool>) {
if self.auto_continue != value {
self.auto_continue = value;
self.dirty = true;
}
}
pub fn set_max_auto_continues(&mut self, value: Option<usize>) {
if self.max_auto_continues != value {
self.max_auto_continues = value;
self.dirty = true;
}
}
pub fn inject_todo_instructions(&self) -> Option<bool> {
self.inject_todo_instructions
}
pub fn continuation_prompt(&self) -> Option<&str> {
self.continuation_prompt.as_deref()
}
pub fn set_inject_todo_instructions(&mut self, value: Option<bool>) {
if self.inject_todo_instructions != value {
self.inject_todo_instructions = value;
self.dirty = true;
}
}
pub fn set_continuation_prompt(&mut self, value: Option<String>) {
if self.continuation_prompt != value {
self.continuation_prompt = value;
self.dirty = true;
}
}
pub fn needs_compression(&self, global_compression_threshold: usize) -> bool {
if self.compressing {
return false;
@@ -548,15 +628,6 @@ impl Session {
let mut messages = self.messages.clone();
if input.continue_output().is_some() {
return messages;
} else if input.regenerate() {
while let Some(last) = messages.last() {
if !last.role.is_user() {
messages.pop();
} else {
break;
}
}
return messages;
}
let mut need_add_msg = true;
let len = messages.len();
+2 -1
View File
@@ -8,6 +8,7 @@ use serde_json::{Value, json};
use std::collections::HashMap;
use std::sync::Arc;
#[derive(Clone)]
pub struct ToolScope {
pub functions: Functions,
pub mcp_runtime: McpRuntime,
@@ -24,7 +25,7 @@ impl Default for ToolScope {
}
}
#[derive(Default)]
#[derive(Default, Clone)]
pub struct McpRuntime {
pub servers: HashMap<String, Arc<ConnectedServer>>,
}
+244
View File
@@ -0,0 +1,244 @@
use crate::utils::warning_text;
use anyhow::{Context, Result, bail};
use dunce::canonicalize;
use inquire::Confirm;
use is_terminal::IsTerminal;
use self_update::Status;
use self_update::backends::github::Update;
use std::fs::OpenOptions;
use std::path::Path;
use std::{env, fs, io, process};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum InstallSource {
Cargo,
Homebrew,
Manual,
}
impl InstallSource {
fn is_package_managed(self) -> bool {
matches!(self, InstallSource::Cargo | InstallSource::Homebrew)
}
fn label(self) -> &'static str {
match self {
InstallSource::Cargo => "Cargo",
InstallSource::Homebrew => "Homebrew",
InstallSource::Manual => "manually-installed",
}
}
}
fn classify_install_path(path: &Path) -> InstallSource {
let components: Vec<&str> = path
.components()
.filter_map(|c| c.as_os_str().to_str())
.collect();
if components
.windows(2)
.any(|w| w[0] == ".cargo" && w[1] == "bin")
{
return InstallSource::Cargo;
}
if components.contains(&"Cellar") {
return InstallSource::Homebrew;
}
let path_str = path.to_string_lossy();
if path_str.starts_with("/opt/homebrew/") || path_str.starts_with("/home/linuxbrew/.linuxbrew/")
{
return InstallSource::Homebrew;
}
InstallSource::Manual
}
fn normalize_version(requested: Option<String>) -> Option<String> {
let raw = requested?;
let trimmed = raw.trim();
if trimmed.is_empty() || trimmed.eq_ignore_ascii_case("latest") {
return None;
}
match trimmed.chars().next() {
Some('v' | 'V') => Some(trimmed.to_string()),
Some(c) if c.is_ascii_digit() => Some(format!("v{trimmed}")),
_ => Some(trimmed.to_string()),
}
}
fn is_dir_writable(dir: &Path) -> bool {
let probe = dir.join(format!(".loki-update-write-test-{}", process::id()));
match OpenOptions::new().write(true).create_new(true).open(&probe) {
Ok(_) => {
let _ = fs::remove_file(&probe);
true
}
Err(_) => false,
}
}
pub fn run_self_update(requested: Option<String>, force: bool) -> Result<()> {
let target_tag = normalize_version(requested);
let exe_path = env::current_exe()
.context("Could not determine the path of the running loki executable")?;
let resolved = canonicalize(&exe_path).unwrap_or_else(|_| exe_path.clone());
let source = classify_install_path(&resolved);
if source.is_package_managed() {
let body = match source {
InstallSource::Homebrew => format!(
"Loki appears to be installed via Homebrew ({}).\n\
Updating in place replaces the binary inside Homebrew's Cellar; `brew` will\n\
then report a version that no longer matches the file on disk, and a later\n\
`brew upgrade`/`brew reinstall` may overwrite it or fail.\n\
The clean way to update is: brew upgrade loki",
exe_path.display()
),
InstallSource::Cargo => format!(
"Loki appears to be installed via `cargo install` ({}).\n\
Updating in place leaves Cargo's records out of sync with the binary on disk.\n\
The clean way to update is: cargo install --locked loki-ai",
exe_path.display()
),
InstallSource::Manual => unreachable!("Manual installs are not package-managed"),
};
println!("{} {body}", warning_text("WARNING:"));
if force {
println!("--force specified; updating anyway.");
} else if io::stdin().is_terminal() {
let proceed = Confirm::new("Update anyway?")
.with_default(false)
.prompt()?;
if !proceed {
println!("Update cancelled.");
return Ok(());
}
} else {
bail!(
"Refusing to update a {} install. Re-run with --force to override.",
source.label()
);
}
}
if let Some(parent) = exe_path.parent()
&& !is_dir_writable(parent)
{
bail!(
"No write permission for '{}'. Re-run with elevated permissions (e.g. sudo), \
or update Loki through your package manager.",
parent.display()
);
}
let interactive = io::stdin().is_terminal();
let mut builder = Update::configure();
builder
.repo_owner("Dark-Alex-17")
.repo_name("loki")
.bin_name("loki")
.current_version(env!("CARGO_PKG_VERSION"))
.no_confirm(true)
.show_download_progress(interactive);
if let Some(tag) = &target_tag {
builder.target_version_tag(tag.as_str());
}
let status = builder
.build()
.context("Failed to configure the self-update")?
.update()
.context("Self-update failed")?;
match status {
Status::UpToDate(version) => {
println!("Loki is already up to date (v{version}).");
}
Status::Updated(version) => {
println!("Loki updated to v{version}. Restart loki to use the new version.");
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
#[test]
fn classify_cargo_install() {
assert_eq!(
classify_install_path(&PathBuf::from("/home/u/.cargo/bin/loki")),
InstallSource::Cargo
);
}
#[test]
fn classify_homebrew_opt_prefix() {
assert_eq!(
classify_install_path(&PathBuf::from("/opt/homebrew/bin/loki")),
InstallSource::Homebrew
);
}
#[test]
fn classify_homebrew_cellar() {
assert_eq!(
classify_install_path(&PathBuf::from("/usr/local/Cellar/loki/0.3.0/bin/loki")),
InstallSource::Homebrew
);
}
#[test]
fn classify_homebrew_linuxbrew() {
assert_eq!(
classify_install_path(&PathBuf::from("/home/linuxbrew/.linuxbrew/bin/loki")),
InstallSource::Homebrew
);
}
#[test]
fn classify_manual_usr_local_bin() {
assert_eq!(
classify_install_path(&PathBuf::from("/usr/local/bin/loki")),
InstallSource::Manual
);
}
#[test]
fn classify_manual_local_bin() {
assert_eq!(
classify_install_path(&PathBuf::from("/home/u/.local/bin/loki")),
InstallSource::Manual
);
}
#[test]
fn normalize_version_latest_and_empty_are_none() {
assert_eq!(normalize_version(None), None);
assert_eq!(normalize_version(Some(String::new())), None);
assert_eq!(normalize_version(Some(" ".to_string())), None);
assert_eq!(normalize_version(Some("latest".to_string())), None);
assert_eq!(normalize_version(Some("LATEST".to_string())), None);
}
#[test]
fn normalize_version_prepends_v_for_bare_semver() {
assert_eq!(
normalize_version(Some("0.4.0".to_string())),
Some("v0.4.0".to_string())
);
assert_eq!(
normalize_version(Some("v0.4.0".to_string())),
Some("v0.4.0".to_string())
);
assert_eq!(
normalize_version(Some(" v0.4.0 ".to_string())),
Some("v0.4.0".to_string())
);
}
}
+58 -8
View File
@@ -51,7 +51,7 @@ enum BinaryType<'a> {
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, AsRefStr)]
enum Language {
pub enum Language {
Bash,
Python,
TypeScript,
@@ -60,7 +60,13 @@ enum Language {
impl From<&String> for Language {
fn from(s: &String) -> Self {
match s.to_lowercase().as_str() {
Language::from_extension(s)
}
}
impl Language {
pub fn from_extension(ext: &str) -> Self {
match ext.to_lowercase().as_str() {
"sh" => Language::Bash,
"py" => Language::Python,
"ts" => Language::TypeScript,
@@ -90,6 +96,17 @@ impl Language {
}
}
impl Language {
pub fn direct_invoker(self) -> Option<(&'static str, &'static [&'static str])> {
match self {
Language::Bash => Some(("bash", &[])),
Language::Python => Some(("python3", &[])),
Language::TypeScript => Some(("npx", &["tsx"])),
Language::Unsupported => None,
}
}
}
fn extract_shebang_runtime(path: &Path) -> Option<String> {
let file = File::open(path).ok()?;
let reader = io::BufReader::new(file);
@@ -192,7 +209,7 @@ pub struct Functions {
}
impl Functions {
pub fn install_builtin_global_tools() -> Result<()> {
pub fn install_builtin_global_tools(force: bool) -> Result<()> {
info!(
"Installing global built-in functions in {}",
paths::functions_dir().display()
@@ -210,14 +227,14 @@ impl Functions {
})?;
let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
let file_path = paths::functions_dir().join(file.as_ref());
let file_extension = file_path
#[cfg_attr(not(unix), expect(unused))]
let is_script = file_path
.extension()
.and_then(OsStr::to_str)
.map(|s| s.to_lowercase());
#[cfg_attr(not(unix), expect(unused))]
let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py"));
.is_some_and(|ext| Language::from_extension(ext) != Language::Unsupported);
if file_path.exists() {
let force_this = force && file.as_ref() != "mcp.json";
if file_path.exists() && !force_this {
debug!(
"Function file already exists, skipping: {}",
file_path.display()
@@ -240,6 +257,22 @@ impl Functions {
Ok(())
}
pub fn install_mcp_config() -> Result<()> {
let file_path = paths::mcp_config_file();
let embedded = FunctionAssets::get("mcp.json")
.ok_or_else(|| anyhow!("Failed to load embedded mcp.json"))?;
let content = unsafe { std::str::from_utf8_unchecked(&embedded.data) };
ensure_parent_exists(&file_path)?;
info!("Reinstalling MCP config file: {}", file_path.display());
let mut config_file = File::create(&file_path)?;
config_file.write_all(content.as_bytes())?;
Ok(())
}
pub fn init(visible_tools: &[String]) -> Result<Self> {
Self::clear_global_functions_bin_dir()?;
@@ -1415,6 +1448,23 @@ mod tests {
assert!(tc.thought_signature.is_none());
}
#[test]
fn direct_invoker_maps_each_language() {
assert_eq!(
Language::Bash.direct_invoker(),
Some(("bash", &[] as &[&str]))
);
assert_eq!(
Language::Python.direct_invoker(),
Some(("python3", &[] as &[&str]))
);
assert_eq!(
Language::TypeScript.direct_invoker(),
Some(("npx", &["tsx"] as &[&str]))
);
assert_eq!(Language::Unsupported.direct_invoker(), None);
}
#[test]
fn toolcall_with_thought_signature() {
let tc = ToolCall::new("t".into(), json!({}), None)
+160 -6
View File
@@ -5,6 +5,7 @@ use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox};
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor};
use crate::utils::{AbortSignal, create_abort_signal};
use crate::graph;
use anyhow::{Context, Result, anyhow, bail};
use chrono::Utc;
use indexmap::IndexMap;
@@ -13,6 +14,8 @@ use parking_lot::RwLock;
use serde_json::{Value, json};
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tokio::time;
use uuid::Uuid;
pub const SUPERVISOR_FUNCTION_PREFIX: &str = "agent__";
@@ -324,12 +327,21 @@ pub async fn handle_supervisor_tool(
}
}
fn run_child_agent(
pub fn run_child_agent(
mut child_ctx: RequestContext,
initial_input: Input,
abort_signal: AbortSignal,
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
Box::pin(async move {
if graph::active_agent_graph_name(&child_ctx).is_some() {
return graph::run_active_agent_graph(
&mut child_ctx,
&initial_input.text(),
abort_signal,
)
.await;
}
let mut accumulated_output = String::new();
let mut input = initial_input;
let app = Arc::clone(&child_ctx.app.config);
@@ -372,6 +384,98 @@ fn run_child_agent(
})
}
/// Spawn an agent synchronously from a graph node and return its accumulated
/// output. This is similar to `handle_spawn` but runs the child agent in the
/// current task (no tokio::spawn, no supervisor handle registration) so the
/// graph executor can sequence agent nodes directly.
pub async fn run_agent_for_graph(
parent_ctx: &mut RequestContext,
agent_name: &str,
prompt: &str,
) -> Result<String> {
let short_uuid = &Uuid::new_v4().to_string()[..8];
let agent_id = format!("graph_agent_{agent_name}_{short_uuid}");
let current_depth = parent_ctx.current_depth + 1;
if let Some(supervisor) = parent_ctx.supervisor.as_ref().cloned() {
let max_depth = supervisor.read().max_depth();
if current_depth > max_depth {
bail!("Max agent depth exceeded ({current_depth}/{max_depth})");
}
}
if !parent_ctx.app.config.function_calling_support {
bail!("Function calling support must be enabled to spawn agents.");
}
let child_inbox = Arc::new(Inbox::new());
parent_ctx.ensure_root_escalation_queue();
let child_abort = create_abort_signal();
let app_config = Arc::clone(&parent_ctx.app.config);
let current_model = parent_ctx.current_model().clone();
let info_flag = parent_ctx.info_flag;
let child_app_state = Arc::new(AppState {
config: Arc::new(app_config.as_ref().clone()),
vault: parent_ctx.app.vault.clone(),
mcp_factory: parent_ctx.app.mcp_factory.clone(),
rag_cache: parent_ctx.app.rag_cache.clone(),
mcp_config: parent_ctx.app.mcp_config.clone(),
mcp_log_path: parent_ctx.app.mcp_log_path.clone(),
mcp_registry: parent_ctx.app.mcp_registry.clone(),
functions: parent_ctx.app.functions.clone(),
});
let agent = Agent::init(
app_config.as_ref(),
child_app_state.as_ref(),
&current_model,
info_flag,
agent_name,
child_abort.clone(),
)
.await?;
let agent_mcp_servers = agent.mcp_server_names().to_vec();
let session = agent.agent_session().map(|v| v.to_string());
let should_init_supervisor = agent.can_spawn_agents();
let agent_max_concurrent = agent.max_concurrent_agents();
let agent_max_depth = agent.max_agent_depth();
let mut child_ctx = RequestContext::new_for_child(
Arc::clone(&child_app_state),
parent_ctx,
current_depth,
Arc::clone(&child_inbox),
agent_id.clone(),
);
child_ctx.rag = agent.rag();
child_ctx.agent = Some(agent);
if should_init_supervisor {
child_ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new(
agent_max_concurrent,
agent_max_depth,
))));
}
if let Some(session) = session {
child_ctx
.use_session(app_config.as_ref(), Some(&session), child_abort.clone())
.await?;
sync_agent_functions_to_ctx(&mut child_ctx)?;
} else {
populate_agent_mcp_runtime(&mut child_ctx, &agent_mcp_servers).await?;
sync_agent_functions_to_ctx(&mut child_ctx)?;
child_ctx.init_agent_shared_variables()?;
}
let input = Input::from_str(&child_ctx, prompt, None);
debug!("Spawning agent '{agent_name}' for graph node as '{agent_id}'");
run_child_agent(child_ctx, input, child_abort).await
}
async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[String]) -> Result<()> {
if !ctx.app.config.mcp_server_support {
return Ok(());
@@ -601,11 +705,25 @@ async fn handle_check(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
match is_finished {
Some(true) => handle_collect(ctx, args).await,
Some(false) => Ok(json!({
Some(false) => {
let mut result = json!({
"status": "pending",
"id": id,
"message": "Agent is still running"
})),
});
if let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending()
{
let summary = queue.pending_summary();
result["pending_escalations"] = json!(summary);
result["message"] = json!(
"Agent is still running. Child agents have pending escalations that need your reply via agent__reply_escalation."
);
}
Ok(result)
}
None => Ok(json!({
"status": "error",
"message": format!("No agent found with id '{id}'")
@@ -619,12 +737,48 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
.and_then(Value::as_str)
.ok_or_else(|| anyhow!("'id' is required"))?;
let handle = {
let supervisor = ctx
.supervisor
.as_ref()
.cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?;
{
let sup = supervisor.read();
if sup.is_finished(id).is_none() {
return Ok(json!({
"status": "error",
"message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.")
}));
}
}
loop {
let is_finished = {
let sup = supervisor.read();
sup.is_finished(id).unwrap_or(false)
};
if is_finished {
break;
}
if let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending()
{
let summary = queue.pending_summary();
return Ok(json!({
"status": "pending",
"id": id,
"message": format!("Agent '{id}' is still running, but child agents have pending escalations that need your reply. Reply via agent__reply_escalation, then call agent__collect again."),
"pending_escalations": summary,
}));
}
time::sleep(Duration::from_millis(200)).await;
}
let handle = {
let mut sup = supervisor.write();
sup.take(id)
};
@@ -649,7 +803,7 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
}
None => Ok(json!({
"status": "error",
"message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.")
"message": format!("Agent '{id}' completed but could not be collected. It may have been collected by another call.")
})),
}
}
@@ -1193,7 +1347,7 @@ mod tests {
let inbox = Arc::new(Inbox::new());
let abort = create_abort_signal();
let join_handle = tokio::spawn(async {
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
time::sleep(Duration::from_secs(60)).await;
Ok(AgentResult {
id: "slow".into(),
agent_name: "test".into(),
+8 -2
View File
@@ -94,8 +94,14 @@ pub fn handle_todo_tool(ctx: &mut RequestContext, cmd_name: &str, args: &Value)
.strip_prefix(TODO_FUNCTION_PREFIX)
.unwrap_or(cmd_name);
if ctx.agent.is_none() {
bail!("No active agent");
if !ctx.app.config.function_calling_support {
bail!("Cannot use todo tools: function calling is disabled.");
}
let auto_config = ctx.auto_continue_config();
if !auto_config.enabled {
bail!(
"Auto-continue is not enabled. Set 'auto_continue: true' in your config to use todo tools."
);
}
match action {
+28 -15
View File
@@ -2,7 +2,7 @@ use super::{FunctionDeclaration, JsonSchema};
use crate::config::RequestContext;
use crate::supervisor::escalation::{EscalationRequest, new_escalation_id};
use anyhow::{Result, anyhow};
use anyhow::{Result, anyhow, bail};
use indexmap::IndexMap;
use inquire::{Confirm, MultiSelect, Select, Text};
use serde_json::{Value, json};
@@ -155,7 +155,10 @@ fn handle_direct_ask(args: &Value) -> Result<Value> {
let mut options = parse_options(args)?;
options.push(CUSTOM_MULTI_CHOICE_ANSWER_OPTION.to_string());
let mut answer = Select::new(question, options).prompt()?;
let mut answer = Select::new(question, options)
.without_filtering()
.with_help_message("↑↓ to move, enter to select")
.prompt()?;
if answer == CUSTOM_MULTI_CHOICE_ANSWER_OPTION {
answer = Text::new("Custom response:").prompt()?
@@ -205,12 +208,11 @@ async fn handle_escalated(ctx: &RequestContext, action: &str, args: &Value) -> R
.ok_or_else(|| anyhow!("'question' is required"))?
.to_string();
let options: Option<Vec<String>> = args.get("options").and_then(Value::as_array).map(|arr| {
arr.iter()
.filter_map(Value::as_str)
.map(String::from)
.collect()
});
let options: Option<Vec<String>> = if args.get("options").is_some() {
Some(parse_options(args)?)
} else {
None
};
let from_agent_id = ctx
.self_agent_id
@@ -262,13 +264,24 @@ async fn handle_escalated(ctx: &RequestContext, action: &str, args: &Value) -> R
}
fn parse_options(args: &Value) -> Result<Vec<String>> {
args.get("options")
.and_then(Value::as_array)
.map(|arr| {
arr.iter()
let raw = args
.get("options")
.ok_or_else(|| anyhow!("'options' is required and must be an array of strings"))?;
let arr: Vec<Value> = match raw {
Value::Array(arr) => arr.clone(),
Value::String(s) => serde_json::from_str::<Vec<Value>>(s).map_err(|_| {
anyhow!(
"'options' was a string but did not parse as a JSON array. \
Pass options as a native JSON array, e.g. [\"yes\", \"no\"]."
)
})?,
_ => bail!("'options' is required and must be an array of strings"),
};
Ok(arr
.iter()
.filter_map(Value::as_str)
.map(String::from)
.collect()
})
.ok_or_else(|| anyhow!("'options' is required and must be an array of strings"))
.collect())
}
+272
View File
@@ -0,0 +1,272 @@
use super::state::StateManager;
use super::structured;
use super::types::AgentNode;
use crate::config::RequestContext;
use crate::function::supervisor::run_agent_for_graph;
use anyhow::{Context, Result};
use serde_json::Value;
use std::time::Duration;
use tokio::time::timeout;
const OUTPUT_KEY: &str = "output";
const DEFAULT_TIMEOUT_SECS: u64 = 300;
pub struct AgentNodeExecutor;
impl AgentNodeExecutor {
pub async fn execute(
node: &AgentNode,
state_manager: &mut StateManager,
parent_ctx: &mut RequestContext,
) -> Result<String> {
let prompt = state_manager
.interpolate(&node.prompt)
.with_context(|| format!("Failed to interpolate prompt for agent '{}'", node.agent))?;
let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS));
let raw = timeout(
timeout_dur,
run_agent_for_graph(parent_ctx, &node.agent, &prompt),
)
.await
.with_context(|| {
format!(
"Agent '{}' timed out after {}s",
node.agent,
timeout_dur.as_secs()
)
})?
.with_context(|| format!("Agent '{}' failed", node.agent))?;
let output_value = match &node.output_schema {
Some(schema) => structured::extract(&raw, schema, parent_ctx)
.await
.with_context(|| {
format!(
"Agent '{}' output failed structured-output extraction",
node.agent
)
})?,
None => Value::String(raw.clone()),
};
apply_state_updates(node, state_manager, &output_value);
Ok(raw)
}
}
fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, output: &Value) {
if node.output_schema.is_some()
&& let Some(obj) = output.as_object()
{
for (k, v) in obj {
state_manager.state_mut().set(k.clone(), v.clone());
}
}
let Some(updates) = &node.state_updates else {
return;
};
let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
state_manager
.state_mut()
.set(OUTPUT_KEY.into(), output.clone());
for (key, template) in updates {
let value = state_manager.interpolate_lenient(template);
state_manager
.state_mut()
.set(key.clone(), Value::String(value));
}
match prev_output {
Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
None => {
state_manager
.state_mut()
.set(OUTPUT_KEY.into(), Value::Null);
}
}
}
#[cfg(test)]
mod tests {
use super::super::types::AgentNode;
use super::*;
use serde_json::json;
use std::collections::HashMap;
fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
let mut map = HashMap::new();
for (k, v) in pairs {
map.insert((*k).into(), v.clone());
}
StateManager::new(map)
}
fn node_with(prompt: &str, updates: Option<HashMap<String, String>>) -> AgentNode {
AgentNode {
agent: "test_agent".into(),
prompt: prompt.into(),
state_updates: updates,
output_schema: None,
timeout: None,
}
}
#[test]
fn state_updates_use_output_placeholder() {
let node = {
let mut u = HashMap::new();
u.insert("findings".into(), "{{output}}".into());
node_with("hi", Some(u))
};
let mut state = manager_with(&[]);
apply_state_updates(&node, &mut state, &json!("agent finished its work"));
assert_eq!(
state.state().get("findings"),
Some(&json!("agent finished its work"))
);
}
#[test]
fn state_updates_can_reference_existing_keys_and_output() {
let node = {
let mut u = HashMap::new();
u.insert("summary".into(), "{{topic}}: {{output}}".into());
node_with("hi", Some(u))
};
let mut state = manager_with(&[("topic", json!("auth"))]);
apply_state_updates(&node, &mut state, &json!("JWT vs sessions"));
assert_eq!(
state.state().get("summary"),
Some(&json!("auth: JWT vs sessions"))
);
}
#[test]
fn output_key_is_cleaned_up_after_state_updates() {
let node = {
let mut u = HashMap::new();
u.insert("findings".into(), "{{output}}".into());
node_with("hi", Some(u))
};
let mut state = manager_with(&[]);
apply_state_updates(&node, &mut state, &json!("anything"));
assert_eq!(state.state().get("output"), Some(&Value::Null));
}
#[test]
fn pre_existing_output_value_is_preserved() {
let node = {
let mut u = HashMap::new();
u.insert("greeting".into(), "{{output}}".into());
node_with("hi", Some(u))
};
let mut state = manager_with(&[("output", json!("preserved"))]);
apply_state_updates(&node, &mut state, &json!("new agent output"));
assert_eq!(
state.state().get("greeting"),
Some(&json!("new agent output"))
);
assert_eq!(state.state().get("output"), Some(&json!("preserved")));
}
#[test]
fn no_state_updates_is_a_noop() {
let node = node_with("hi", None);
let mut state = manager_with(&[("k", json!("v"))]);
apply_state_updates(&node, &mut state, &json!("ignored"));
assert_eq!(state.state().get("k"), Some(&json!("v")));
assert!(state.state().get("output").is_none());
}
#[test]
fn interpolate_lenient_on_state_updates_handles_missing_keys() {
let node = {
let mut u = HashMap::new();
u.insert("decorated".into(), "[{{missing}}] {{output}}".into());
node_with("hi", Some(u))
};
let mut state = manager_with(&[]);
apply_state_updates(&node, &mut state, &json!("DATA"));
assert_eq!(state.state().get("decorated"), Some(&json!("[] DATA")));
}
fn node_with_schema(
prompt: &str,
updates: Option<HashMap<String, String>>,
schema: Value,
) -> AgentNode {
let mut n = node_with(prompt, updates);
n.output_schema = Some(schema);
n
}
#[test]
fn output_schema_auto_merges_top_level_keys() {
let node = node_with_schema("hi", None, json!({"type": "object"}));
let mut state = manager_with(&[]);
let output = json!({"goal": "do X", "summary": "details"});
apply_state_updates(&node, &mut state, &output);
assert_eq!(state.state().get("goal"), Some(&json!("do X")));
assert_eq!(state.state().get("summary"), Some(&json!("details")));
}
#[test]
fn output_schema_preserves_nested_value_types() {
let node = node_with_schema("hi", None, json!({"type": "object"}));
let mut state = manager_with(&[]);
let output = json!({
"tags": ["a", "b"],
"config": { "key": "value" },
"count": 42
});
apply_state_updates(&node, &mut state, &output);
assert_eq!(state.state().get("tags"), Some(&json!(["a", "b"])));
assert_eq!(state.state().get("config"), Some(&json!({"key": "value"})));
assert_eq!(state.state().get("count"), Some(&json!(42)));
}
#[test]
fn output_schema_explicit_state_updates_override_auto_merge() {
let mut u = HashMap::new();
u.insert("goal".into(), "renamed-{{output.goal}}".into());
let node = node_with_schema("hi", Some(u), json!({"type": "object"}));
let mut state = manager_with(&[]);
let output = json!({"goal": "do X"});
apply_state_updates(&node, &mut state, &output);
assert_eq!(state.state().get("goal"), Some(&json!("renamed-do X")));
}
#[test]
fn no_schema_does_not_auto_merge() {
let node = node_with("hi", None);
let mut state = manager_with(&[]);
let output = json!({"goal": "do X"});
apply_state_updates(&node, &mut state, &output);
assert!(state.state().get("goal").is_none());
}
}
+47
View File
@@ -0,0 +1,47 @@
use super::{GraphExecutor, GraphParser, agent_has_graph};
use crate::config::RequestContext;
use crate::config::paths;
use crate::utils::AbortSignal;
use anyhow::{Context, Result, anyhow};
use log::info;
use serde_json::Value;
pub fn active_agent_graph_name(ctx: &RequestContext) -> Option<String> {
let name = ctx.agent.as_ref()?.name().to_string();
agent_has_graph(&name).then_some(name)
}
pub async fn run_active_agent_graph(
ctx: &mut RequestContext,
prompt: &str,
abort_signal: AbortSignal,
) -> Result<String> {
let agent_name =
active_agent_graph_name(ctx).ok_or_else(|| anyhow!("Active agent has no graph.yaml"))?;
info!("Agent '{agent_name}' has graph.yaml; routing to graph executor");
let agent_dir = paths::agent_data_dir(&agent_name);
let graph_path = paths::agent_graph_file(&agent_name);
let parser = GraphParser::new(&agent_dir);
let mut graph = parser
.load_from_file(&graph_path)
.with_context(|| format!("Failed to load graph.yaml for agent '{agent_name}'"))?;
graph
.initial_state
.insert("initial_prompt".into(), Value::String(prompt.to_string()));
let executor = GraphExecutor::new(graph, agent_dir);
let output = executor
.execute(ctx, abort_signal)
.await
.with_context(|| format!("Graph execution failed for agent '{agent_name}'"))?;
if let Some(supervisor) = ctx.supervisor.clone() {
supervisor.read().cancel_all();
}
Ok(output)
}
+796
View File
@@ -0,0 +1,796 @@
use super::agent::AgentNodeExecutor;
use super::llm::{LlmExecutionOutcome, LlmNodeExecutor};
use super::logging::{GraphLogger, narrate_node_complete, narrate_node_failed};
use super::map::MapNodeExecutor;
use super::rag::RagNodeExecutor;
use super::script::ScriptExecutor;
use super::staging::BranchWrites;
use super::state::StateManager;
use super::types::{EndNode, Graph, Node, NodeType};
use super::user_interaction::{ApprovalNodeExecutor, InputNodeExecutor};
use super::validator::{AgentValidationContext, GraphValidator};
use crate::config::{RenderMode, RequestContext};
use crate::utils::AbortSignal;
use anyhow::{Context, Result, anyhow, bail};
use futures_util::future::join_all;
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;
pub struct GraphExecutor {
graph: Graph,
base_dir: PathBuf,
}
impl GraphExecutor {
pub fn new(graph: Graph, base_dir: impl Into<PathBuf>) -> Self {
Self {
graph,
base_dir: base_dir.into(),
}
}
pub async fn execute(
self,
ctx: &mut RequestContext,
abort_signal: AbortSignal,
) -> Result<String> {
let is_nested = ctx.current_depth > 0;
let mut logger = GraphLogger::with_visibility(
&self.graph.name,
self.graph.settings.log_state_snapshots,
is_nested,
);
let result = self.run(&mut logger, ctx, abort_signal).await;
if let Err(e) = &result {
logger.graph_error(e);
}
result
}
async fn run(
self,
logger: &mut GraphLogger,
ctx: &mut RequestContext,
abort_signal: AbortSignal,
) -> Result<String> {
let GraphExecutor { graph, base_dir } = self;
if graph.settings.validate_before_run {
let mut validator = GraphValidator::new(&base_dir);
if let Some(agent) = &ctx.agent {
validator = validator.with_agent_context(AgentValidationContext::from_agent(
agent,
Arc::clone(&ctx.app.config),
));
}
let result = validator.validate(&graph);
for w in &result.warnings {
logger.validation_warning(w.node_id.as_deref(), &w.message);
}
result.into_result()?;
}
let mut state = StateManager::new(graph.initial_state.clone());
let agent_envs = ctx
.agent
.as_ref()
.map(|a| a.variable_envs())
.unwrap_or_default();
let script_executor = ScriptExecutor::new(&base_dir).with_envs(agent_envs);
let max_iterations = graph.settings.max_loop_iterations;
let graph_timeout = graph.settings.timeout.map(Duration::from_secs);
let max_concurrency = graph.settings.max_concurrency;
let graph = Arc::new(graph);
let start = Instant::now();
let mut frontier: HashSet<String> = HashSet::from([graph.start.clone()]);
logger.graph_start(&graph.start, graph.nodes.len());
loop {
if frontier.is_empty() {
bail!(
"Graph '{}' frontier emptied without reaching an End node",
graph.name
);
}
if abort_signal.aborted() {
bail!(
"Graph '{}' aborted before super-step with frontier {:?}",
graph.name,
sorted_frontier(&frontier)
);
}
if let Some(t) = graph_timeout
&& start.elapsed() > t
{
bail!(
"Graph '{}' timed out after {}s before super-step with frontier {:?}",
graph.name,
t.as_secs(),
sorted_frontier(&frontier)
);
}
// Loop-count and visit tracking on live state, BEFORE forking.
// This counts every entry to a node toward max_loop_iterations
// regardless of how many parallel branches converged on it.
for node_id in &frontier {
state.state_mut().visit_node(node_id);
let visits = state.state().loop_count(node_id);
if visits > max_iterations {
bail!(
"Node '{}' visited {} times (max_loop_iterations={}). \
Possible infinite loop.",
node_id,
visits,
max_iterations
);
}
}
for node_id in &frontier {
let node = graph.get_node(node_id).ok_or_else(|| {
anyhow!("Node '{}' not found in graph '{}'", node_id, graph.name)
})?;
let visits = state.state().loop_count(node_id);
logger.node_entry(node, visits);
}
let snapshot_label = if frontier.len() == 1 {
frontier.iter().next().cloned().unwrap_or_default()
} else {
format!("super-step {{{}}}", sorted_frontier(&frontier).join(","))
};
logger.state_snapshot(&snapshot_label, &state);
let snapshot = state.read_snapshot();
let semaphore = Arc::new(Semaphore::new(max_concurrency));
let frontier_size = frontier.len();
let in_super_step = frontier_size > 1;
let silent = logger.silent();
if in_super_step {
let mut branches = sorted_frontier(&frontier);
branches.sort();
logger.super_step_start(&branches);
}
let mut branch_tasks = Vec::with_capacity(frontier_size);
for node_id in &frontier {
let node = graph
.get_node(node_id)
.ok_or_else(|| {
anyhow!("Node '{}' not found in graph '{}'", node_id, graph.name)
})?
.clone();
logger.node_start(&node, in_super_step);
let branch_state = state.fork_for_branch_state();
let mut branch_ctx = ctx.fork_for_branch();
if in_super_step {
branch_ctx.render_mode = RenderMode::Silent;
}
let script_exec_clone = script_executor.clone();
let graph_clone = Arc::clone(&graph);
let current = node_id.clone();
let sem_clone = semaphore.clone();
let abort_clone = abort_signal.clone();
let task = tokio::spawn(async move {
let _permit = sem_clone
.acquire()
.await
.expect("semaphore should not be closed");
if abort_clone.aborted() {
narrate_node_failed(
silent,
&node,
Duration::default(),
"aborted",
in_super_step,
);
return (
current.clone(),
branch_state,
Err(anyhow!("branch aborted")),
Duration::default(),
);
}
let node_start = Instant::now();
let mut state = branch_state;
let mut ctx = branch_ctx;
let step_ctx = StepContext {
graph: graph_clone.as_ref(),
script_executor: &script_exec_clone,
max_concurrency,
abort_signal: &abort_clone,
};
let result = step(&node, &mut state, &mut ctx, &step_ctx, &current).await;
let elapsed = node_start.elapsed();
match &result {
Ok(StepResult::Continue(targets)) => {
let route = if targets.is_empty() {
None
} else {
Some(targets.join(", "))
};
narrate_node_complete(
silent,
&node,
elapsed,
route.as_deref(),
in_super_step,
);
}
Ok(StepResult::End(_)) => {
narrate_node_complete(
silent,
&node,
elapsed,
Some("END"),
in_super_step,
);
}
Err(e) => {
narrate_node_failed(
silent,
&node,
elapsed,
&e.to_string(),
in_super_step,
);
}
}
(current, state, result, elapsed)
});
branch_tasks.push(task);
}
let joined = join_all(branch_tasks).await;
let mut branch_writes: Vec<BranchWrites> = Vec::new();
let mut next_frontier: HashSet<String> = HashSet::new();
let mut end_results: Vec<(String, StateManager, String)> = Vec::new();
for join_result in joined {
let (node_id, branch_state, step_result, elapsed) =
join_result.map_err(|e| anyhow!("Branch task panicked: {e}"))?;
logger.record_timing(&node_id, elapsed);
let step_outcome = step_result.with_context(|| format!("at node '{node_id}'"))?;
match step_outcome {
StepResult::Continue(targets) => {
for target in &targets {
logger.routing(&node_id, target);
}
let diff = branch_state.diff_against(snapshot.as_ref());
branch_writes.push(BranchWrites {
node_id: node_id.clone(),
invocation_index: 0,
writes: diff,
});
next_frontier.extend(targets);
}
StepResult::End(output) => {
end_results.push((node_id.clone(), branch_state, output));
}
}
}
if end_results.len() > 1 {
let mut ids: Vec<String> =
end_results.iter().map(|(id, _, _)| id.clone()).collect();
ids.sort();
bail!(
"super-step ended with multiple End targets ({}). \
Fan-out branches must converge at a join node before \
terminating. To fix: route all parallel branches to a \
single shared next-node, then terminate from there.",
ids.join(", ")
);
}
// Sort by (node_id, invocation_index) so non-commutative reducers
// like Concat/Merge produce deterministic output across runs.
branch_writes.sort_by(|a, b| {
a.node_id
.cmp(&b.node_id)
.then(a.invocation_index.cmp(&b.invocation_index))
});
state.apply_branch_writes(branch_writes, &graph.reducers)?;
if let Some((node_id, end_state, output)) = end_results.into_iter().next() {
let diff = end_state.diff_against(snapshot.as_ref());
state.apply_branch_writes(
vec![BranchWrites {
node_id: node_id.clone(),
invocation_index: 0,
writes: diff,
}],
&graph.reducers,
)?;
logger.graph_complete(&node_id, start.elapsed());
return Ok(output);
}
if in_super_step {
logger.super_step_end(&sorted_frontier(&next_frontier));
}
frontier = next_frontier;
}
}
}
fn sorted_frontier(frontier: &HashSet<String>) -> Vec<String> {
let mut v: Vec<String> = frontier.iter().cloned().collect();
v.sort();
v
}
pub(super) struct StepContext<'a> {
pub graph: &'a Graph,
pub script_executor: &'a ScriptExecutor,
pub max_concurrency: usize,
pub abort_signal: &'a AbortSignal,
}
impl StepContext<'_> {
pub fn graph_name(&self) -> &str {
&self.graph.name
}
}
enum StepResult {
// The set of next-node ids the executor should add to the next super-step's
// frontier. A `Vec` of length 1 for sequential routing (default) and the
// full target list for fan-out (`next: [a, b, ...]`). Dynamic single-route
// decisions (script `_next`, approval routes, LLM/RAG fallback) always emit
// a single-element vec.
Continue(Vec<String>),
End(String),
}
async fn step(
node: &Node,
state: &mut StateManager,
ctx: &mut RequestContext,
step_ctx: &StepContext<'_>,
current: &str,
) -> Result<StepResult> {
match &node.node_type {
NodeType::Agent(agent_node) => {
AgentNodeExecutor::execute(agent_node, state, ctx).await?;
let targets = static_next_targets(node, current, "agent")?;
Ok(StepResult::Continue(targets))
}
NodeType::Script(script_node) => {
let dynamic = match step_ctx.script_executor.execute(script_node, state).await {
Ok(n) => n,
Err(e) => {
if let Some(fallback) = &script_node.fallback {
warn!(
"[graph:{}] script '{}' failed, routing to fallback '{}': {}",
step_ctx.graph_name(),
current,
fallback,
e
);
return Ok(StepResult::Continue(vec![fallback.clone()]));
}
return Err(e);
}
};
let targets = match dynamic {
Some(n) => vec![n],
None => static_next_targets(node, current, "script")?,
};
Ok(StepResult::Continue(targets))
}
NodeType::Approval(approval_node) => {
let next = ApprovalNodeExecutor::execute(approval_node, state, ctx).await?;
Ok(StepResult::Continue(vec![next]))
}
NodeType::Input(input_node) => {
let next_id = first_next_target(node);
let next = InputNodeExecutor::execute(input_node, next_id, state, ctx).await?;
Ok(StepResult::Continue(vec![next]))
}
NodeType::Llm(llm_node) => {
let outcome = LlmNodeExecutor::execute(llm_node, state, ctx).await?;
let targets = match outcome {
LlmExecutionOutcome::Continue => static_next_targets(node, current, "llm")?,
LlmExecutionOutcome::FellBack(target) => vec![target],
};
Ok(StepResult::Continue(targets))
}
NodeType::Rag(rag_node) => {
RagNodeExecutor::execute(rag_node, current, state, ctx).await?;
let targets = static_next_targets(node, current, "rag")?;
Ok(StepResult::Continue(targets))
}
NodeType::End(end_node) => Ok(StepResult::End(resolve_end_output(end_node, state))),
NodeType::Map(map_node) => {
let targets = static_next_targets(node, current, "map")?;
MapNodeExecutor::execute(map_node, state, ctx, step_ctx, current).await?;
Ok(StepResult::Continue(targets))
}
}
}
fn static_next_targets(node: &Node, current: &str, kind: &str) -> Result<Vec<String>> {
node.next
.as_ref()
.map(|t| t.as_slice().to_vec())
.ok_or_else(|| anyhow!("{kind} node '{current}' has no `next` and is not an end node"))
}
fn first_next_target(node: &Node) -> Option<&str> {
node.next
.as_ref()
.and_then(|t| t.as_slice().first().map(|s| s.as_str()))
}
fn resolve_end_output(end_node: &EndNode, state: &mut StateManager) -> String {
apply_simple_state_updates(end_node.state_updates.as_ref(), state);
state.interpolate_lenient(&end_node.output)
}
fn apply_simple_state_updates(updates: Option<&HashMap<String, String>>, state: &mut StateManager) {
let Some(updates) = updates else {
return;
};
for (key, template) in updates {
let value = state.interpolate_lenient(template);
state.state_mut().set(key.clone(), Value::String(value));
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn state_with(pairs: &[(&str, Value)]) -> StateManager {
let mut map = HashMap::new();
for (k, v) in pairs {
map.insert((*k).into(), v.clone());
}
StateManager::new(map)
}
fn end_node(output: &str, updates: Option<HashMap<String, String>>) -> EndNode {
EndNode {
output: output.into(),
state_updates: updates,
}
}
#[test]
fn resolve_end_output_interpolates_template_against_state() {
let mut state = state_with(&[("name", json!("alice"))]);
let node = end_node("done: {{name}}", None);
assert_eq!(resolve_end_output(&node, &mut state), "done: alice");
}
#[test]
fn resolve_end_output_applies_state_updates_before_interpolation() {
let mut updates = HashMap::new();
updates.insert("summary".into(), "completed for {{user}}".into());
let node = end_node("RESULT: {{summary}}", Some(updates));
let mut state = state_with(&[("user", json!("bob"))]);
assert_eq!(
resolve_end_output(&node, &mut state),
"RESULT: completed for bob"
);
assert_eq!(
state.state().get("summary"),
Some(&json!("completed for bob"))
);
}
#[test]
fn resolve_end_output_with_empty_template_returns_empty_string() {
let mut state = state_with(&[]);
let node = end_node("", None);
assert_eq!(resolve_end_output(&node, &mut state), "");
}
#[test]
fn resolve_end_output_lenient_on_missing_keys() {
let mut state = state_with(&[]);
let node = end_node("hello {{unknown}}!", None);
assert_eq!(resolve_end_output(&node, &mut state), "hello !");
}
#[test]
fn apply_simple_state_updates_does_nothing_when_none() {
let mut state = state_with(&[("k", json!("v"))]);
apply_simple_state_updates(None, &mut state);
assert_eq!(state.state().get("k"), Some(&json!("v")));
}
#[test]
fn apply_simple_state_updates_overwrites_existing_values() {
let mut updates = HashMap::new();
updates.insert("k".into(), "new-{{k}}".into());
let mut state = state_with(&[("k", json!("old"))]);
apply_simple_state_updates(Some(&updates), &mut state);
assert_eq!(state.state().get("k"), Some(&json!("new-old")));
}
}
#[cfg(test)]
mod integration_tests {
use super::*;
use crate::config::{AppState, WorkingMode};
use crate::utils::{create_abort_signal, temp_file};
use std::fs;
fn cmd_available(name: &str) -> bool {
which::which(name).is_ok()
}
struct TestWorkspace {
dir: PathBuf,
}
impl TestWorkspace {
fn new() -> Self {
let dir = temp_file("-graph-integration-", "");
fs::create_dir_all(&dir).unwrap();
Self { dir }
}
fn write_script(&self, name: &str, contents: &str) {
fs::write(self.dir.join(name), contents).unwrap();
}
}
impl Drop for TestWorkspace {
fn drop(&mut self) {
let _ = fs::remove_dir_all(&self.dir);
}
}
fn make_ctx() -> RequestContext {
RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd)
}
#[tokio::test]
async fn static_fan_out_merges_branch_writes_via_append_reducer() {
if !cmd_available("bash") {
eprintln!("skipping: bash not available");
return;
}
let ws = TestWorkspace::new();
ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n");
ws.write_script(
"worker_a.sh",
"#!/bin/bash\necho '{\"results\": \"alpha\"}'\n",
);
ws.write_script(
"worker_b.sh",
"#!/bin/bash\necho '{\"results\": \"beta\"}'\n",
);
let yaml = r#"
name: static_fan_out_test
start: dispatcher
reducers:
results: append
nodes:
dispatcher:
type: script
script: dispatcher.sh
state_updates: {}
next: [worker_a, worker_b]
worker_a:
type: script
script: worker_a.sh
state_updates: {}
next: join
worker_b:
type: script
script: worker_b.sh
state_updates: {}
next: join
join:
type: end
output: "{{results}}"
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
let mut ctx = make_ctx();
let abort = create_abort_signal();
let result = GraphExecutor::new(graph, &ws.dir)
.execute(&mut ctx, abort)
.await
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
let parsed: Value = serde_json::from_str(&result)
.unwrap_or_else(|_| panic!("expected JSON array, got: {result}"));
let arr = parsed.as_array().expect("results should be an array");
assert_eq!(arr.len(), 2, "expected 2 elements, got: {result}");
let strs: Vec<&str> = arr.iter().filter_map(|v| v.as_str()).collect();
assert!(strs.contains(&"alpha"), "missing 'alpha' in {strs:?}");
assert!(strs.contains(&"beta"), "missing 'beta' in {strs:?}");
}
#[tokio::test]
async fn map_over_list_collects_outputs_in_input_order() {
if !cmd_available("python3") {
eprintln!("skipping: python3 not available");
return;
}
let ws = TestWorkspace::new();
ws.write_script(
"doubler.py",
r#"#!/usr/bin/env python3
import os, json
state = json.loads(os.environ.get("GRAPH_STATE", "{}"))
val = state["item"]
print(json.dumps({"output": val * 2}))
"#,
);
let yaml = r#"
name: map_input_order_test
start: fan_out
initial_state:
items: [1, 2, 3, 4, 5]
nodes:
fan_out:
type: map
over: "{{items}}"
as: item
branch: doubler
collect_into: doubled
next: done
doubler:
type: script
script: doubler.py
state_updates: {}
done:
type: end
output: "{{doubled}}"
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
let mut ctx = make_ctx();
let abort = create_abort_signal();
let result = GraphExecutor::new(graph, &ws.dir)
.execute(&mut ctx, abort)
.await
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
let parsed: Value = serde_json::from_str(&result)
.unwrap_or_else(|_| panic!("expected JSON array, got: {result}"));
let arr = parsed.as_array().expect("doubled should be an array");
let nums: Vec<i64> = arr
.iter()
.map(|v| v.as_i64().expect("each item should be int"))
.collect();
assert_eq!(
nums,
vec![2, 4, 6, 8, 10],
"map outputs should be in input order, not finish order"
);
}
#[tokio::test]
async fn parallel_branch_error_aborts_super_step() {
if !cmd_available("bash") {
eprintln!("skipping: bash not available");
return;
}
let ws = TestWorkspace::new();
ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n");
ws.write_script(
"worker_ok.sh",
"#!/bin/bash\necho '{\"results\": \"ok\"}'\n",
);
ws.write_script(
"worker_fail.sh",
"#!/bin/bash\necho 'simulated failure' >&2\nexit 1\n",
);
let yaml = r#"
name: branch_error_test
start: dispatcher
reducers:
results: append
nodes:
dispatcher:
type: script
script: dispatcher.sh
state_updates: {}
next: [worker_ok, worker_fail]
worker_ok:
type: script
script: worker_ok.sh
state_updates: {}
next: join
worker_fail:
type: script
script: worker_fail.sh
state_updates: {}
next: join
join:
type: end
output: "{{results}}"
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
let mut ctx = make_ctx();
let abort = create_abort_signal();
let result = GraphExecutor::new(graph, &ws.dir)
.execute(&mut ctx, abort)
.await;
assert!(result.is_err(), "expected branch error to propagate");
let err = format!("{:#}", result.unwrap_err());
assert!(
err.contains("worker_fail"),
"error should mention failing node: {err}"
);
}
#[tokio::test]
async fn multi_end_in_super_step_is_rejected() {
if !cmd_available("bash") {
eprintln!("skipping: bash not available");
return;
}
let ws = TestWorkspace::new();
ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n");
let yaml = r#"
name: multi_end_test
start: dispatcher
nodes:
dispatcher:
type: script
script: dispatcher.sh
state_updates: {}
next: [end_a, end_b]
end_a:
type: end
output: "from a"
end_b:
type: end
output: "from b"
"#;
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
let mut ctx = make_ctx();
let abort = create_abort_signal();
let result = GraphExecutor::new(graph, &ws.dir)
.execute(&mut ctx, abort)
.await;
assert!(result.is_err(), "expected multi-End to be rejected");
let err = format!("{:#}", result.unwrap_err());
assert!(
err.contains("multiple End targets"),
"error should explain multi-End cause: {err}"
);
assert!(
err.contains("end_a") && err.contains("end_b"),
"error should list both End nodes: {err}"
);
}
}
+616
View File
@@ -0,0 +1,616 @@
use super::state::StateManager;
use super::structured;
use super::types::LlmNode;
use crate::client::{Model, ModelType, call_chat_completions};
use crate::config::{Input, RequestContext, Role, RoleLike};
use crate::utils::create_abort_signal;
use anyhow::{Context, Error, Result, anyhow, bail};
use serde_json::Value;
use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::timeout;
const OUTPUT_KEY: &str = "output";
#[derive(Debug, PartialEq, Eq)]
pub(super) enum LlmExecutionOutcome {
Continue,
FellBack(String),
}
pub struct LlmNodeExecutor;
impl LlmNodeExecutor {
pub(super) async fn execute(
node: &LlmNode,
state_manager: &mut StateManager,
parent_ctx: &mut RequestContext,
) -> Result<LlmExecutionOutcome> {
let result = run(node, state_manager, parent_ctx).await;
let (output, failure_reason) = match result {
Ok(raw) => match &node.output_schema {
Some(schema) => match structured::extract(&raw, schema, parent_ctx).await {
Ok(value) => (value, None),
Err(e) => {
warn!("llm node structured extraction failed: {e}");
(
Value::String(format!("LLM node structured-extraction failed: {e}")),
Some(format!("structured-extraction failed: {e}")),
)
}
},
None => (Value::String(raw), None),
},
Err(e) => {
warn!("llm node failed: {e}");
(
Value::String(format!("LLM node failed: {e}")),
Some(format!("LLM call failed: {e:#}")),
)
}
};
apply_state_updates_with_output(node, state_manager, &output);
outcome_from(failure_reason.as_deref(), node.fallback.as_deref())
}
}
fn outcome_from(
failure_reason: Option<&str>,
fallback: Option<&str>,
) -> Result<LlmExecutionOutcome> {
match (failure_reason, fallback) {
(None, _) => Ok(LlmExecutionOutcome::Continue),
(Some(_), Some(fb)) => Ok(LlmExecutionOutcome::FellBack(fb.to_string())),
(Some(reason), None) => bail!(
"LLM node failed and no fallback declared: {reason}. \
Add a `fallback:` route on the node to route on failure, \
or fix the underlying error."
),
}
}
async fn run(
node: &LlmNode,
state_manager: &mut StateManager,
parent_ctx: &mut RequestContext,
) -> Result<String> {
let mut instructions: Option<String> = match &node.instructions {
Some(s) => Some(
state_manager
.interpolate(s)
.context("Failed to interpolate llm node instructions")?,
),
None => None,
};
let mut prompt = state_manager
.interpolate(&node.prompt)
.context("Failed to interpolate llm node prompt")?;
if let Some(schema) = &node.output_schema {
let hint = format_schema_hint(schema);
match instructions.as_mut() {
Some(s) => {
s.push_str("\n\n");
s.push_str(&hint);
}
None => {
prompt.push_str("\n\n");
prompt.push_str(&hint);
}
}
}
let (regular_tools, mcp_servers) = categorize_tools(node.tools.as_deref());
validate_tools_subset(&regular_tools, &mcp_servers, parent_ctx)?;
let role = build_inline_role(
node,
instructions.as_deref(),
&regular_tools,
&mcp_servers,
parent_ctx,
)?;
let saved_role = parent_ctx.role.clone();
parent_ctx.role = Some(role);
let result = match node.timeout {
Some(secs) => match timeout(
Duration::from_secs(secs),
run_with_retries(node, &prompt, parent_ctx),
)
.await
{
Ok(r) => r,
Err(_) => Err(anyhow!("llm node timed out after {secs}s")),
},
None => run_with_retries(node, &prompt, parent_ctx).await,
};
parent_ctx.role = saved_role;
result
}
async fn run_with_retries(
node: &LlmNode,
prompt: &str,
ctx: &mut RequestContext,
) -> Result<String> {
let mut last_err: Option<Error> = None;
for attempt in 1..=node.max_attempts {
match run_chat_loop(node, prompt, ctx).await {
Ok(out) => return Ok(out),
Err(e) if is_transient(&e) && attempt < node.max_attempts => {
warn!("llm node attempt {attempt} failed (transient): {e}; retrying");
last_err = Some(e);
}
Err(e) => return Err(e),
}
}
Err(last_err.unwrap_or_else(|| anyhow!("llm node exhausted retries")))
}
async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) -> Result<String> {
let abort = create_abort_signal();
let app_cfg = Arc::clone(&ctx.app.config);
let role_for_input = ctx.role.clone();
let mut input = Input::from_str(ctx, prompt, role_for_input);
let mut accumulated = String::new();
for turn in 0..node.max_iterations {
let client = input.create_client()?;
ctx.before_chat_completion(&input)?;
let (output, tool_results) =
call_chat_completions(&input, false, false, client.as_ref(), ctx, abort.clone())
.await?;
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
if !output.is_empty() {
if !accumulated.is_empty() {
accumulated.push('\n');
}
accumulated.push_str(&output);
}
if tool_results.is_empty() {
return Ok(accumulated);
}
if turn + 1 == node.max_iterations {
bail!(
"llm node hit max_iterations ({}) before LLM concluded",
node.max_iterations
);
}
input = input.merge_tool_results(output, tool_results);
}
bail!("llm node ended without producing output")
}
fn build_inline_role(
node: &LlmNode,
instructions: Option<&str>,
regular_tools: &[String],
mcp_servers: &[String],
parent_ctx: &RequestContext,
) -> Result<Role> {
let mut role = Role::new("llm_node", instructions.unwrap_or(""));
let model = match &node.model {
Some(model_id) => {
Model::retrieve_model(parent_ctx.app.config.as_ref(), model_id, ModelType::Chat)
.with_context(|| format!("Unknown model '{model_id}' on llm node"))?
}
None => parent_ctx.current_model().clone(),
};
role.set_model(model);
if let Some(t) = node.temperature {
role.set_temperature(Some(t));
}
if let Some(p) = node.top_p {
role.set_top_p(Some(p));
}
if node.tools.as_deref().unwrap_or_default().is_empty() {
role.set_enabled_tools(Some(String::new()));
role.set_enabled_mcp_servers(Some(String::new()));
} else {
if !regular_tools.is_empty() {
role.set_enabled_tools(Some(regular_tools.join(",")));
} else {
role.set_enabled_tools(Some(String::new()));
}
if !mcp_servers.is_empty() {
role.set_enabled_mcp_servers(Some(mcp_servers.join(",")));
} else {
role.set_enabled_mcp_servers(Some(String::new()));
}
}
Ok(role)
}
fn categorize_tools(entries: Option<&[String]>) -> (Vec<String>, Vec<String>) {
let mut regular = Vec::new();
let mut mcp = Vec::new();
let Some(entries) = entries else {
return (regular, mcp);
};
for e in entries {
if let Some(server) = e.strip_prefix("mcp:") {
mcp.push(server.to_string());
} else {
regular.push(e.clone());
}
}
(regular, mcp)
}
fn validate_tools_subset(
regular: &[String],
mcp_servers: &[String],
parent_ctx: &RequestContext,
) -> Result<()> {
let agent = parent_ctx
.agent
.as_ref()
.ok_or_else(|| anyhow!("llm node requires an active agent"))?;
if !regular.is_empty() {
let known: HashSet<&str> = agent
.functions()
.declarations()
.iter()
.map(|d| d.name.as_str())
.collect();
for name in regular {
if !known.contains(name.as_str()) {
let mut avail: Vec<&str> = known.iter().copied().collect();
avail.sort();
bail!(
"llm node references unknown tool '{name}'. Agent '{}' provides: {}",
agent.name(),
avail.join(", ")
);
}
}
}
if !mcp_servers.is_empty() {
let known: HashSet<&str> = agent
.mcp_server_names()
.iter()
.map(|s| s.as_str())
.collect();
for server in mcp_servers {
if !known.contains(server.as_str()) {
let mut avail: Vec<&str> = known.iter().copied().collect();
avail.sort();
bail!(
"llm node references unknown MCP server 'mcp:{server}'. \
Agent '{}' has MCP servers: [{}]",
agent.name(),
avail.join(", ")
);
}
}
}
Ok(())
}
fn is_transient(err: &Error) -> bool {
let s = format!("{err:#}");
s.contains("timed out")
|| s.contains("rate limit")
|| s.contains("429")
|| s.contains("Connection reset")
|| s.contains("Connection refused")
|| s.contains("produced no output")
}
fn apply_state_updates_with_output(
node: &LlmNode,
state_manager: &mut StateManager,
output: &Value,
) {
if node.output_schema.is_some()
&& let Some(obj) = output.as_object()
{
for (k, v) in obj {
state_manager.state_mut().set(k.clone(), v.clone());
}
}
let Some(updates) = &node.state_updates else {
return;
};
let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
state_manager
.state_mut()
.set(OUTPUT_KEY.into(), output.clone());
for (key, template) in updates {
let value = state_manager.interpolate_lenient(template);
state_manager
.state_mut()
.set(key.clone(), Value::String(value));
}
match prev_output {
Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
None => {
state_manager
.state_mut()
.set(OUTPUT_KEY.into(), Value::Null);
}
}
}
fn format_schema_hint(schema: &Value) -> String {
let schema_json = serde_json::to_string_pretty(schema).unwrap_or_else(|_| schema.to_string());
format!(
"Respond with a JSON object that matches this schema. Output ONLY the JSON \
object with no surrounding prose or markdown fences.\n\nSchema:\n{schema_json}"
)
}
#[cfg(test)]
mod tests {
use super::super::types::*;
use super::*;
use serde_json::json;
use std::collections::HashMap;
fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
let mut map = HashMap::new();
for (k, v) in pairs {
map.insert((*k).into(), v.clone());
}
StateManager::new(map)
}
fn node_with(updates: Option<HashMap<String, String>>) -> LlmNode {
LlmNode {
instructions: Some("sys".into()),
prompt: "user".into(),
tools: None,
model: None,
temperature: None,
top_p: None,
fallback: None,
max_attempts: 1,
max_iterations: 10,
state_updates: updates,
output_schema: None,
timeout: None,
}
}
#[test]
fn state_updates_expose_output_during_evaluation() {
let mut u = HashMap::new();
u.insert("response".into(), "{{output}}".into());
let node = node_with(Some(u));
let mut state = manager_with(&[]);
apply_state_updates_with_output(&node, &mut state, &json!("the answer"));
assert_eq!(state.state().get("response"), Some(&json!("the answer")));
}
#[test]
fn state_updates_can_mix_existing_keys_with_output() {
let mut u = HashMap::new();
u.insert("summary".into(), "{{topic}}: {{output}}".into());
let node = node_with(Some(u));
let mut state = manager_with(&[("topic", json!("LOINC"))]);
apply_state_updates_with_output(&node, &mut state, &json!("abc"));
assert_eq!(state.state().get("summary"), Some(&json!("LOINC: abc")));
}
#[test]
fn output_key_is_cleared_after_state_updates() {
let mut u = HashMap::new();
u.insert("k".into(), "{{output}}".into());
let node = node_with(Some(u));
let mut state = manager_with(&[]);
apply_state_updates_with_output(&node, &mut state, &json!("anything"));
assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null)));
}
#[test]
fn pre_existing_output_value_is_restored() {
let mut u = HashMap::new();
u.insert("greeting".into(), "{{output}}".into());
let node = node_with(Some(u));
let mut state = manager_with(&[("output", json!("preserved"))]);
apply_state_updates_with_output(&node, &mut state, &json!("new"));
assert_eq!(state.state().get("greeting"), Some(&json!("new")));
assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!("preserved")));
}
#[test]
fn no_state_updates_is_a_noop() {
let node = node_with(None);
let mut state = manager_with(&[("k", json!("v"))]);
apply_state_updates_with_output(&node, &mut state, &json!("x"));
assert_eq!(state.state().get("k"), Some(&json!("v")));
assert!(state.state().get(OUTPUT_KEY).is_none());
}
#[test]
fn outcome_from_success_is_continue() {
assert_eq!(
outcome_from(None, Some("fb")).unwrap(),
LlmExecutionOutcome::Continue
);
assert_eq!(
outcome_from(None, None).unwrap(),
LlmExecutionOutcome::Continue
);
}
#[test]
fn outcome_from_failure_with_fallback_is_fell_back() {
assert_eq!(
outcome_from(Some("HTTP 404"), Some("fb")).unwrap(),
LlmExecutionOutcome::FellBack("fb".to_string())
);
}
#[test]
fn outcome_from_failure_without_fallback_propagates_error() {
let err = outcome_from(Some("HTTP 404"), None).unwrap_err();
let msg = format!("{err:#}");
assert!(msg.contains("no fallback declared"), "got: {msg}");
assert!(msg.contains("HTTP 404"), "got: {msg}");
}
fn node_with_schema(updates: Option<HashMap<String, String>>, schema: Value) -> LlmNode {
let mut n = node_with(updates);
n.output_schema = Some(schema);
n
}
#[test]
fn output_schema_auto_merges_top_level_keys() {
let node = node_with_schema(None, json!({"type": "object"}));
let mut state = manager_with(&[]);
let output = json!({"goal": "do X", "summary": "details"});
apply_state_updates_with_output(&node, &mut state, &output);
assert_eq!(state.state().get("goal"), Some(&json!("do X")));
assert_eq!(state.state().get("summary"), Some(&json!("details")));
}
#[test]
fn output_schema_preserves_nested_value_types() {
let node = node_with_schema(None, json!({"type": "object"}));
let mut state = manager_with(&[]);
let output = json!({
"tags": ["a", "b"],
"config": { "key": "value" },
"count": 42
});
apply_state_updates_with_output(&node, &mut state, &output);
assert_eq!(state.state().get("tags"), Some(&json!(["a", "b"])));
assert_eq!(state.state().get("config"), Some(&json!({"key": "value"})));
assert_eq!(state.state().get("count"), Some(&json!(42)));
}
#[test]
fn output_schema_explicit_state_updates_override_auto_merge() {
let mut u = HashMap::new();
u.insert("goal".into(), "renamed-{{output.goal}}".into());
let node = node_with_schema(Some(u), json!({"type": "object"}));
let mut state = manager_with(&[]);
let output = json!({"goal": "do X"});
apply_state_updates_with_output(&node, &mut state, &output);
assert_eq!(state.state().get("goal"), Some(&json!("renamed-do X")));
}
#[test]
fn output_schema_skips_auto_merge_for_non_object() {
let node = node_with_schema(None, json!({"type": "array"}));
let mut state = manager_with(&[]);
let output = json!([1, 2, 3]);
apply_state_updates_with_output(&node, &mut state, &output);
assert!(state.state().get("0").is_none());
assert!(state.state().get(OUTPUT_KEY).is_none());
}
#[test]
fn no_schema_does_not_auto_merge() {
let node = node_with(None);
let mut state = manager_with(&[]);
let output = json!({"goal": "do X"});
apply_state_updates_with_output(&node, &mut state, &output);
assert!(state.state().get("goal").is_none());
}
#[test]
fn format_schema_hint_includes_schema_and_instruction() {
let schema = json!({"type": "object", "properties": {"goal": {"type": "string"}}});
let hint = format_schema_hint(&schema);
assert!(hint.contains("Schema:"));
assert!(hint.contains("\"goal\""));
assert!(hint.contains("JSON"));
assert!(hint.contains("ONLY"));
}
#[test]
fn categorize_tools_splits_mcp_and_regular() {
let entries = vec![
"read_query".to_string(),
"mcp:pubmed-search".to_string(),
"web_search_loki".to_string(),
"mcp:github".to_string(),
];
let (regular, mcp) = categorize_tools(Some(&entries));
assert_eq!(regular, vec!["read_query", "web_search_loki"]);
assert_eq!(mcp, vec!["pubmed-search", "github"]);
}
#[test]
fn categorize_tools_with_none_returns_empty() {
let (regular, mcp) = categorize_tools(None);
assert!(regular.is_empty());
assert!(mcp.is_empty());
}
#[test]
fn categorize_tools_with_empty_returns_empty() {
let (regular, mcp) = categorize_tools(Some(&[]));
assert!(regular.is_empty());
assert!(mcp.is_empty());
}
#[test]
fn is_transient_matches_expected_signatures() {
assert!(is_transient(&anyhow!("request timed out after 30s")));
assert!(is_transient(&anyhow!("rate limit reached")));
assert!(is_transient(&anyhow!("429 too many requests")));
assert!(is_transient(&anyhow!("Connection reset by peer")));
assert!(is_transient(&anyhow!("Connection refused")));
assert!(is_transient(&anyhow!("llm produced no output")));
}
#[test]
fn is_transient_rejects_non_transient_errors() {
assert!(!is_transient(&anyhow!("Unknown model 'foo'")));
assert!(!is_transient(&anyhow!(
"llm node references unknown tool 'bad'"
)));
assert!(!is_transient(&anyhow!("hit max_iterations")));
assert!(!is_transient(&anyhow!("authentication failed")));
}
}
+312
View File
@@ -0,0 +1,312 @@
use super::state::StateManager;
use super::types::{Node, NodeType};
use crate::utils::dimmed_text;
use chrono::Local;
use indexmap::IndexMap;
use std::cmp::Reverse;
use std::time::Duration;
fn ts() -> String {
Local::now().format("%H:%M:%S").to_string()
}
fn fmt_secs(elapsed: Duration) -> String {
let secs = elapsed.as_secs_f64();
if secs < 1.0 {
format!("{}ms", elapsed.as_millis())
} else {
format!("{secs:.2}s")
}
}
#[derive(Debug, Clone, Default)]
struct NodeTiming {
count: usize,
total: Duration,
max: Duration,
}
impl NodeTiming {
fn record(&mut self, elapsed: Duration) {
self.count += 1;
self.total += elapsed;
if elapsed > self.max {
self.max = elapsed;
}
}
}
pub struct GraphLogger {
graph_name: String,
log_state_snapshots: bool,
silent: bool,
timings: IndexMap<String, NodeTiming>,
}
impl GraphLogger {
pub fn with_visibility(graph_name: &str, log_state_snapshots: bool, silent: bool) -> Self {
Self {
graph_name: graph_name.to_string(),
log_state_snapshots,
silent,
timings: IndexMap::new(),
}
}
pub fn graph_start(&self, start_node: &str, node_count: usize) {
info!(
"[graph:{}] start at '{}' ({} nodes)",
self.graph_name, start_node, node_count
);
if !self.silent {
eprintln!(
"{}",
dimmed_text(&format!(
"▸ graph: {} (start: {start_node})",
self.graph_name
))
);
}
}
pub fn graph_complete(&self, end_node: &str, elapsed: Duration) {
info!(
"[graph:{}] end '{}' (elapsed {:?})",
self.graph_name, end_node, elapsed
);
if !self.silent {
eprintln!(
"{}",
dimmed_text(&format!("▸ graph done in {:.2}s", elapsed.as_secs_f64()))
);
}
self.log_performance_summary();
}
pub fn graph_error(&self, error: &anyhow::Error) {
error!("[graph:{}] execution failed: {error:#}", self.graph_name);
}
pub fn node_entry(&self, node: &Node, visit: usize) {
debug!(
"[graph:{}] entering '{}' (visit {visit})",
self.graph_name, node.id
);
}
pub fn silent(&self) -> bool {
self.silent
}
pub fn node_start(&self, node: &Node, in_super_step: bool) {
narrate_node_start(self.silent, node, in_super_step);
}
pub fn super_step_start(&self, branches: &[String]) {
if self.silent {
return;
}
eprintln!(
"{}",
dimmed_text(&format!(
"▸ {} super-step start: {}",
ts(),
branches.join(", ")
))
);
}
pub fn super_step_end(&self, targets: &[String]) {
if self.silent {
return;
}
let route = if targets.is_empty() {
String::new()
} else {
format!(" -> {}", targets.join(", "))
};
eprintln!(
"{}",
dimmed_text(&format!("{} super-step end{route}", ts()))
);
}
pub fn record_timing(&mut self, node_id: &str, elapsed: Duration) {
self.timings
.entry(node_id.to_string())
.or_default()
.record(elapsed);
}
pub fn routing(&self, from: &str, to: &str) {
debug!("[graph:{}] {from} -> {to}", self.graph_name);
}
pub fn validation_warning(&self, node_id: Option<&str>, message: &str) {
match node_id {
Some(id) => warn!("[graph:{}] [{id}] {message}", self.graph_name),
None => warn!("[graph:{}] {message}", self.graph_name),
}
}
pub fn state_snapshot(&self, node_id: &str, state: &StateManager) {
if !self.log_state_snapshots {
return;
}
let snapshot = state.snapshot();
let mut keys: Vec<&str> = snapshot.keys().map(String::as_str).collect();
keys.sort_unstable();
debug!(
"[graph:{}] [{node_id}] state: {} bytes, keys={:?}",
self.graph_name,
state.size_bytes(),
keys
);
trace!(
"[graph:{}] [{node_id}] full state: {:?}",
self.graph_name, snapshot
);
}
fn log_performance_summary(&self) {
if self.timings.is_empty() {
return;
}
let mut rows: Vec<(&String, &NodeTiming)> = self.timings.iter().collect();
rows.sort_by_key(|b| Reverse(b.1.total));
info!(
"[graph:{}] performance summary (slowest first):",
self.graph_name
);
for (node_id, t) in rows {
let avg = t.total / t.count.max(1) as u32;
info!(
"[graph:{}] {node_id}: {} visit(s), total {}ms, avg {}ms, max {}ms",
self.graph_name,
t.count,
t.total.as_millis(),
avg.as_millis(),
t.max.as_millis(),
);
}
}
}
pub fn narrate_node_start(silent: bool, node: &Node, in_super_step: bool) {
if silent {
return;
}
let indent = if in_super_step { " " } else { "" };
let label = node_type_label(node);
eprintln!(
"{}",
dimmed_text(&format!("{} {indent}{} ({label}) start", ts(), node.id))
);
}
pub fn narrate_node_complete(
silent: bool,
node: &Node,
elapsed: Duration,
next_target: Option<&str>,
in_super_step: bool,
) {
if silent {
return;
}
let indent = if in_super_step { " " } else { "" };
let label = node_type_label(node);
let dur = fmt_secs(elapsed);
let route = next_target.map(|t| format!(" -> {t}")).unwrap_or_default();
eprintln!(
"{}",
dimmed_text(&format!(
"▸ {} {indent}{} ({label}) done in {dur}{route}",
ts(),
node.id
))
);
}
pub fn narrate_node_failed(
silent: bool,
node: &Node,
elapsed: Duration,
err: &str,
in_super_step: bool,
) {
if silent {
return;
}
let indent = if in_super_step { " " } else { "" };
let label = node_type_label(node);
let dur = fmt_secs(elapsed);
let excerpt: String = err.chars().take(120).collect();
eprintln!(
"{}",
dimmed_text(&format!(
"▸ {} {indent}{} ({label}) FAILED in {dur} -- {excerpt}",
ts(),
node.id
))
);
}
pub(super) fn node_type_label(node: &Node) -> &'static str {
match &node.node_type {
NodeType::Agent(_) => "agent",
NodeType::Script(_) => "script",
NodeType::Approval(_) => "approval",
NodeType::Input(_) => "input",
NodeType::Llm(_) => "llm",
NodeType::Rag(_) => "rag",
NodeType::End(_) => "end",
NodeType::Map(_) => "map",
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn records_and_aggregates_node_timings() {
let mut logger = GraphLogger::with_visibility("g", false, false);
logger.record_timing("a", Duration::from_millis(100));
logger.record_timing("a", Duration::from_millis(300));
logger.record_timing("b", Duration::from_millis(50));
let a = logger.timings.get("a").unwrap();
assert_eq!(a.count, 2);
assert_eq!(a.total, Duration::from_millis(400));
assert_eq!(a.max, Duration::from_millis(300));
let b = logger.timings.get("b").unwrap();
assert_eq!(b.count, 1);
assert_eq!(b.total, Duration::from_millis(50));
}
#[test]
fn node_timing_max_tracks_largest() {
let mut t = NodeTiming::default();
t.record(Duration::from_millis(10));
t.record(Duration::from_millis(80));
t.record(Duration::from_millis(40));
assert_eq!(t.max, Duration::from_millis(80));
assert_eq!(t.count, 3);
assert_eq!(t.total, Duration::from_millis(130));
}
#[test]
fn new_logger_has_no_timings() {
let logger = GraphLogger::with_visibility("g", true, false);
assert!(logger.timings.is_empty());
assert!(logger.log_state_snapshots);
}
}
+152
View File
@@ -0,0 +1,152 @@
use super::agent::AgentNodeExecutor;
use super::executor::StepContext;
use super::llm::LlmNodeExecutor;
use super::rag::RagNodeExecutor;
use super::state::StateManager;
use super::types::{MapNode, NodeType};
use crate::config::{RenderMode, RequestContext};
use crate::graph::type_name;
use anyhow::{Context, Result, anyhow};
use futures_util::future::join_all;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Semaphore;
pub(super) struct MapNodeExecutor;
impl MapNodeExecutor {
pub(super) async fn execute(
node: &MapNode,
state: &mut StateManager,
ctx: &mut RequestContext,
step_ctx: &StepContext<'_>,
node_id: &str,
) -> Result<()> {
let over_value = state
.interpolate_raw(&node.over)
.with_context(|| format!("map node '{node_id}': evaluating `over` template"))?;
let items = over_value.as_array().ok_or_else(|| {
anyhow!(
"map node '{}': `over` template '{}' must resolve to an array, got {}",
node_id,
node.over,
type_name(&over_value)
)
})?;
let items = items.clone();
let branch_node = step_ctx
.graph
.get_node(&node.branch)
.ok_or_else(|| {
anyhow!(
"map node '{node_id}': branch '{}' not found in graph",
node.branch
)
})?
.clone();
let max_conc = node
.max_concurrency
.unwrap_or(step_ctx.max_concurrency)
.max(1);
let semaphore = Arc::new(Semaphore::new(max_conc));
let mut sub_tasks = Vec::with_capacity(items.len());
for (idx, item) in items.iter().enumerate() {
let item = item.clone();
let as_name = node.as_name.clone();
let branch_clone = branch_node.clone();
let mut sub_state = state.fork_for_branch_state();
let mut sub_ctx = ctx.fork_for_branch();
sub_ctx.render_mode = RenderMode::Silent;
let script_clone = step_ctx.script_executor.clone();
let sub_branch_id = node.branch.clone();
let sem = semaphore.clone();
let abort = step_ctx.abort_signal.clone();
sub_state.state_mut().set(as_name, item);
let task = tokio::spawn(async move {
let _permit = sem
.acquire()
.await
.expect("map semaphore should not be closed");
if abort.aborted() {
return (
idx,
sub_state,
Err(anyhow!("map sub-branch [{idx}] aborted")),
);
}
let mut state = sub_state;
let mut ctx = sub_ctx;
let exec_result: Result<()> = match &branch_clone.node_type {
NodeType::Llm(n) => LlmNodeExecutor::execute(n, &mut state, &mut ctx)
.await
.map(|_| ()),
NodeType::Agent(n) => AgentNodeExecutor::execute(n, &mut state, &mut ctx)
.await
.map(|_| ()),
NodeType::Rag(n) => {
RagNodeExecutor::execute(n, &sub_branch_id, &mut state, &mut ctx).await
}
NodeType::Script(n) => script_clone.execute(n, &mut state).await.map(|_| ()),
_ => Err(anyhow!(
"map branch '{}' has type that cannot run inside a map \
(validator should have caught this; internal error)",
branch_clone.id
)),
};
(idx, state, exec_result)
});
sub_tasks.push(task);
}
let joined = join_all(sub_tasks).await;
// Collect outputs keyed by input index so order is preserved regardless of finish order.
let mut outputs: HashMap<usize, Value> = HashMap::new();
for join_result in joined {
let (idx, sub_state, exec_result) =
join_result.map_err(|e| anyhow!("map sub-branch panicked: {e}"))?;
exec_result
.with_context(|| format!("map node '{node_id}': sub-branch [{idx}] failed"))?;
let output_value = sub_state
.state()
.get(&node.output_key)
.cloned()
.ok_or_else(|| {
anyhow!(
"map node '{node_id}': sub-branch [{idx}] did not write \
`output_key` '{}'",
node.output_key
)
})?;
outputs.insert(idx, output_value);
}
let mut collected = Vec::with_capacity(items.len());
for idx in 0..items.len() {
let value = outputs.remove(&idx).ok_or_else(|| {
anyhow!(
"map node '{node_id}': internal error: missing result for sub-branch [{idx}]"
)
})?;
collected.push(value);
}
state
.state_mut()
.set(node.collect_into.clone(), Value::Array(collected));
Ok(())
}
}
+39
View File
@@ -0,0 +1,39 @@
pub mod agent;
pub mod dispatch;
pub mod executor;
pub mod llm;
pub mod logging;
pub mod map;
pub mod parser;
pub mod rag;
pub mod reducer;
pub mod script;
pub mod staging;
pub mod state;
pub mod structured;
pub mod types;
pub mod user_interaction;
pub mod validator;
pub use dispatch::{active_agent_graph_name, run_active_agent_graph};
pub use executor::GraphExecutor;
pub use parser::{GraphParser, agent_has_graph};
use serde_json::Value;
pub use types::{Graph, NodeType};
pub const GRAPH_SCHEMA_VERSION: &str = "1.0";
pub const DEFAULT_MAX_LOOP_ITERATIONS: usize = 100;
pub const MAX_STATE_SIZE_BYTES: usize = 32 * 1024;
pub(in crate::graph) fn type_name(value: &Value) -> &'static str {
match value {
Value::Null => "null",
Value::Bool(_) => "bool",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
+462
View File
@@ -0,0 +1,462 @@
use super::types::Graph;
use crate::config::paths;
use anyhow::{Context, Error, Result, anyhow, bail};
use std::fs::read_to_string;
use std::path::{Path, PathBuf};
const SUPPORTED_VERSIONS: &[&str] = &["1.0"];
pub struct GraphParser {
base_dir: PathBuf,
}
impl GraphParser {
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
Self {
base_dir: base_dir.into(),
}
}
pub fn load_from_file(&self, path: impl AsRef<Path>) -> Result<Graph> {
let path = path.as_ref();
let full_path = if path.is_absolute() {
path.to_path_buf()
} else {
self.base_dir.join(path)
};
let contents = read_to_string(&full_path)
.with_context(|| format!("Failed to read graph file at '{}'", full_path.display()))?;
self.load_from_string(&contents)
.with_context(|| format!("Failed to parse graph file at '{}'", full_path.display()))
}
pub fn load_from_string(&self, yaml: &str) -> Result<Graph> {
let mut graph: Graph = serde_yaml::from_str(yaml).map_err(enhance_yaml_error)?;
validate_schema_version(&graph.version)?;
for (key, node) in &mut graph.nodes {
if node.id.is_empty() {
node.id = key.clone();
} else if &node.id != key {
bail!(
"Node ID mismatch: key '{}' does not match node.id '{}'",
key,
node.id
);
}
}
validate_structure(&graph)?;
Ok(graph)
}
}
fn validate_schema_version(version: &str) -> Result<()> {
if !SUPPORTED_VERSIONS.contains(&version) {
bail!(
"Unsupported graph schema version '{}'. Supported versions: {}",
version,
SUPPORTED_VERSIONS.join(", ")
);
}
Ok(())
}
fn validate_structure(graph: &Graph) -> Result<()> {
if graph.name.is_empty() {
bail!("Graph must have a non-empty 'name' field");
}
if graph.nodes.is_empty() {
bail!("Graph '{}' has no nodes defined", graph.name);
}
if !graph.has_node(&graph.start) {
bail!(
"Start node '{}' not found in graph '{}'. Available nodes: {}",
graph.start,
graph.name,
graph.node_ids().join(", ")
);
}
Ok(())
}
fn enhance_yaml_error(error: serde_yaml::Error) -> Error {
let msg = error.to_string();
let hint = if msg.contains("missing field") {
"\n\nHint: Check that all required fields are present.\n\
Top-level required fields: `name`, `start`, `nodes`.\n\
Each node requires `type` plus that type's fields:\n\
- agent: `agent`, `prompt`\n\
- script: `script`\n\
- approval: `question`, `options`, `routes`, `on_other`\n\
- input: `question`\n\
- llm: `prompt`\n\
- rag: `documents`\n\
- end: (no required fields)"
} else if msg.contains("unknown field") || msg.contains("unknown variant") {
"\n\nHint: Check for typos in field names or `type:` values.\n\
Valid node types: agent, script, approval, input, llm, rag, end."
} else if msg.contains("invalid type") {
"\n\nHint: Check that field values have the correct type.\n\
- Strings should be quoted if they contain special characters\n\
- Numbers should not be quoted\n\
- Lists use YAML array syntax (- item1)\n\
- Maps use YAML object syntax (key: value)"
} else {
""
};
anyhow!("YAML parsing error: {}{}", msg, hint)
}
pub fn agent_has_graph(agent_name: &str) -> bool {
paths::agent_graph_file(agent_name).exists()
}
#[cfg(test)]
mod tests {
use super::super::GRAPH_SCHEMA_VERSION;
use super::super::types::NodeType;
use super::*;
use indoc::formatdoc;
use std::fs::File;
use std::io::Write;
use std::{env, fs, process};
fn parser() -> GraphParser {
GraphParser::new(env::current_dir().unwrap())
}
#[test]
fn parses_a_simple_graph() {
let yaml = formatdoc! {r#"
name: simple_graph
version: "1.0"
start: node1
nodes:
node1:
id: node1
type: agent
agent: test_agent
prompt: "Hello world"
next: node2
node2:
id: node2
type: end
output: done
"#};
let graph = parser().load_from_string(&yaml).unwrap();
assert_eq!(graph.name, "simple_graph");
assert_eq!(graph.start, "node1");
assert_eq!(graph.nodes.len(), 2);
assert_eq!(
graph.nodes.get("node1").unwrap().next_target(),
Some("node2")
);
}
#[test]
fn auto_fills_node_ids_from_keys() {
let yaml = formatdoc! {r#"
name: auto_id_graph
version: "1.0"
start: node1
nodes:
node1:
type: agent
agent: test_agent
prompt: Test
next: node2
node2:
type: end
output: done
"#};
let graph = parser().load_from_string(&yaml).unwrap();
assert_eq!(graph.nodes.get("node1").unwrap().id, "node1");
assert_eq!(graph.nodes.get("node2").unwrap().id, "node2");
}
#[test]
fn rejects_missing_start_node() {
let yaml = formatdoc! {r#"
name: bad_graph
version: "1.0"
start: nonexistent
nodes:
node1:
type: end
"#};
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
assert!(
err.contains("Start node 'nonexistent' not found"),
"got: {err}"
);
}
#[test]
fn rejects_empty_graph_name() {
let yaml = formatdoc! {r#"
name: ""
version: "1.0"
start: node1
nodes:
node1:
type: end
"#};
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
assert!(err.contains("non-empty 'name'"), "got: {err}");
}
#[test]
fn rejects_no_nodes() {
let yaml = formatdoc! {r#"
name: empty_graph
version: "1.0"
start: node1
nodes: {}
"#, "{}"};
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
assert!(err.contains("no nodes defined"), "got: {err}");
}
#[test]
fn rejects_unsupported_version() {
let yaml = formatdoc! {r#"
name: future_graph
version: "2.0"
start: node1
nodes:
node1:
type: end
"#};
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
assert!(
err.contains("Unsupported graph schema version"),
"got: {err}"
);
}
#[test]
fn rejects_node_id_mismatch() {
let yaml = formatdoc! {r#"
name: mismatch_graph
version: "1.0"
start: node1
nodes:
node1:
id: different_id
type: end
"#};
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
assert!(err.contains("Node ID mismatch"), "got: {err}");
}
#[test]
fn parses_approval_node_with_routes() {
let yaml = formatdoc! {r#"
name: approval_graph
version: "1.0"
start: approval1
nodes:
approval1:
type: approval
question: "Proceed with deployment?"
options:
- "Yes"
- "No"
routes:
"Yes": deploy
"No": cancel
on_other: cancel
deploy:
type: end
cancel:
type: end
"#};
let graph = parser().load_from_string(&yaml).unwrap();
let approval = graph.nodes.get("approval1").unwrap();
match &approval.node_type {
NodeType::Approval(a) => {
assert_eq!(a.options.len(), 2);
assert_eq!(a.routes.len(), 2);
assert_eq!(a.routes.get("Yes").map(|s| s.as_str()), Some("deploy"));
}
_ => panic!("expected approval node"),
}
}
#[test]
fn parses_settings_overrides() {
let yaml = formatdoc! {r#"
name: settings_graph
version: "1.0"
start: node1
settings:
max_loop_iterations: 50
timeout: 300
log_state_snapshots: false
nodes:
node1:
type: end
"#};
let graph = parser().load_from_string(&yaml).unwrap();
assert_eq!(graph.settings.max_loop_iterations, 50);
assert_eq!(graph.settings.timeout, Some(300));
assert!(!graph.settings.log_state_snapshots);
assert!(graph.settings.validate_before_run);
}
#[test]
fn parses_initial_state() {
let yaml = formatdoc! {r#"
name: state_graph
version: "1.0"
start: node1
initial_state:
user_name: "Alice"
count: 42
enabled: true
nodes:
node1:
type: end
"#};
let graph = parser().load_from_string(&yaml).unwrap();
assert_eq!(graph.initial_state.len(), 3);
assert_eq!(graph.initial_state.get("user_name").unwrap(), "Alice");
assert_eq!(
graph.initial_state.get("count").unwrap(),
&serde_json::json!(42)
);
assert_eq!(
graph.initial_state.get("enabled").unwrap(),
&serde_json::json!(true)
);
}
#[test]
fn uses_default_version_when_absent() {
let yaml = formatdoc! {r#"
name: no_version
start: node1
nodes:
node1:
type: end
"#};
let graph = parser().load_from_string(&yaml).unwrap();
assert_eq!(graph.version, GRAPH_SCHEMA_VERSION);
}
#[test]
fn rejects_unknown_node_type_with_hint() {
let yaml = formatdoc! {r#"
name: bad_type
version: "1.0"
start: node1
nodes:
node1:
type: nonsense
"#};
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
assert!(
err.contains("Valid node types") || err.contains("unknown variant"),
"got: {err}"
);
}
#[test]
fn rejects_malformed_yaml() {
let yaml = "name: bad\n bad: indent\nstart: a";
let result = parser().load_from_string(yaml);
assert!(result.is_err());
}
#[test]
fn missing_required_fields_have_a_hint() {
let yaml = formatdoc! {r#"
name: missing_start
version: "1.0"
nodes:
node1:
type: end
"#};
let err = parser().load_from_string(&yaml).unwrap_err().to_string();
assert!(err.contains("Hint"), "got: {err}");
}
#[test]
fn load_from_file_reads_disk() {
let dir = env::temp_dir();
let path = dir.join(format!("loki_graph_parser_test_{}.yaml", process::id()));
let yaml = formatdoc! {r#"
name: disk_graph
version: "1.0"
start: only
nodes:
only:
type: end
output: ok
"#};
{
let mut f = File::create(&path).unwrap();
f.write_all(yaml.as_bytes()).unwrap();
}
let graph = GraphParser::new(dir).load_from_file(&path).unwrap();
assert_eq!(graph.name, "disk_graph");
let _ = fs::remove_file(&path);
}
#[test]
fn load_from_file_errors_on_missing_path() {
let err = parser()
.load_from_file("/definitely/not/a/real/path/to_any_graph.yaml")
.unwrap_err()
.to_string();
assert!(err.contains("Failed to read graph file"), "got: {err}");
}
#[test]
fn agent_has_graph_false_for_unknown_agent() {
assert!(!agent_has_graph("__nonexistent_agent_for_test__"));
}
}
+129
View File
@@ -0,0 +1,129 @@
use super::state::StateManager;
use super::types::RagNode;
use crate::config::RequestContext;
use crate::utils::create_abort_signal;
use anyhow::{Context, Result, anyhow};
use serde_json::{Map, Value};
use std::time::Duration;
use tokio::time::timeout;
const OUTPUT_KEY: &str = "output";
const DEFAULT_QUERY: &str = "{{initial_prompt}}";
const DEFAULT_RAG_TIMEOUT_SECS: u64 = 120;
pub struct RagNodeExecutor;
impl RagNodeExecutor {
pub(super) async fn execute(
node: &RagNode,
node_id: &str,
state_manager: &mut StateManager,
ctx: &mut RequestContext,
) -> Result<()> {
let query_template = node.query.as_deref().unwrap_or(DEFAULT_QUERY);
let query = state_manager
.interpolate(query_template)
.context("Failed to interpolate rag node query")?;
let rag = ctx
.agent
.as_ref()
.and_then(|a| a.graph_rag(node_id))
.ok_or_else(|| anyhow!("rag node '{node_id}' has no initialized knowledge base"))?;
let top_k = node.top_k.unwrap_or_else(|| rag.configured_top_k());
let rerank = rag.configured_reranker();
let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS));
let abort = create_abort_signal();
let (context, sources_str, _ids) =
timeout(timeout_dur, rag.search(&query, top_k, rerank, abort))
.await
.with_context(|| {
format!(
"rag node '{node_id}' timed out after {}s",
timeout_dur.as_secs()
)
})?
.with_context(|| format!("rag node '{node_id}' retrieval failed"))?;
let output = build_rag_output(context, &sources_str);
apply_state_updates(node, state_manager, &output);
Ok(())
}
}
/// Assemble the `{{output}}` value as `{ "context": <ctx>, "sources": [...] }`.
fn build_rag_output(context: String, sources_str: &str) -> Value {
let sources: Vec<Value> = sources_str
.lines()
.map(|line| line.trim().trim_start_matches("- ").trim())
.filter(|s| !s.is_empty())
.map(|s| Value::String(s.to_string()))
.collect();
let mut obj = Map::new();
obj.insert("context".into(), Value::String(context));
obj.insert("sources".into(), Value::Array(sources));
Value::Object(obj)
}
fn apply_state_updates(node: &RagNode, state_manager: &mut StateManager, output: &Value) {
let Some(updates) = &node.state_updates else {
return;
};
let prev_output = state_manager.state().get(OUTPUT_KEY).cloned();
state_manager
.state_mut()
.set(OUTPUT_KEY.into(), output.clone());
for (key, template) in updates {
let value = state_manager.interpolate_lenient(template);
state_manager
.state_mut()
.set(key.clone(), Value::String(value));
}
match prev_output {
Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v),
None => state_manager
.state_mut()
.set(OUTPUT_KEY.into(), Value::Null),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn build_rag_output_splits_bullet_sources_into_array() {
let out = build_rag_output("ctx".into(), "- a.md\n- https://x.com/spec");
assert_eq!(out["context"], json!("ctx"));
assert_eq!(out["sources"], json!(["a.md", "https://x.com/spec"]));
}
#[test]
fn build_rag_output_handles_empty_sources() {
let out = build_rag_output("ctx".into(), "");
assert_eq!(out["sources"], json!([]));
}
#[test]
fn build_rag_output_ignores_blank_lines() {
let out = build_rag_output("c".into(), "- a\n\n- b\n");
assert_eq!(out["sources"], json!(["a", "b"]));
}
#[test]
fn build_rag_output_tolerates_unprefixed_lines() {
let out = build_rag_output("c".into(), "plain/path");
assert_eq!(out["sources"], json!(["plain/path"]));
}
}
+389
View File
@@ -0,0 +1,389 @@
use super::types::Reducer;
use crate::graph::type_name;
use anyhow::{Result, bail};
use serde_json::{Number, Value};
pub fn apply(reducer: Reducer, current: Option<&Value>, incoming: Value) -> Result<Value> {
match reducer {
Reducer::Append => apply_append(current, incoming),
Reducer::Extend => apply_extend(current, incoming),
Reducer::Concat => apply_concat(current, incoming),
Reducer::Sum => apply_sum(current, incoming),
Reducer::Max => apply_max(current, incoming),
Reducer::Min => apply_min(current, incoming),
Reducer::Merge => apply_merge(current, incoming),
Reducer::Overwrite => Ok(incoming),
}
}
fn apply_append(current: Option<&Value>, incoming: Value) -> Result<Value> {
let mut arr = match current {
None => Vec::new(),
Some(Value::Array(a)) => a.clone(),
Some(other) => bail!(
"reducer 'append' requires an array (or absent) for the current value, got {}",
type_name(other)
),
};
arr.push(incoming);
Ok(Value::Array(arr))
}
fn apply_extend(current: Option<&Value>, incoming: Value) -> Result<Value> {
let mut arr = match current {
None => Vec::new(),
Some(Value::Array(a)) => a.clone(),
Some(other) => bail!(
"reducer 'extend' requires an array (or absent) for the current value, got {}",
type_name(other)
),
};
match incoming {
Value::Array(items) => arr.extend(items),
other => bail!(
"reducer 'extend' requires an array for the incoming value, got {}",
type_name(&other)
),
}
Ok(Value::Array(arr))
}
fn apply_concat(current: Option<&Value>, incoming: Value) -> Result<Value> {
let incoming_str = match incoming {
Value::String(s) => s,
other => bail!(
"reducer 'concat' requires a string for the incoming value, got {}",
type_name(&other)
),
};
let result = match current {
None => incoming_str,
Some(Value::String(c)) => {
if c.is_empty() {
incoming_str
} else {
format!("{c}\n{incoming_str}")
}
}
Some(other) => bail!(
"reducer 'concat' requires a string (or absent) for the current value, got {}",
type_name(other)
),
};
Ok(Value::String(result))
}
fn apply_sum(current: Option<&Value>, incoming: Value) -> Result<Value> {
let i = number_or_error(&incoming, "sum", "incoming")?;
let c = match current {
None => 0.0,
Some(value) => number_or_error(value, "sum", "current")?,
};
Ok(json_number(c + i))
}
fn apply_max(current: Option<&Value>, incoming: Value) -> Result<Value> {
let i = number_or_error(&incoming, "max", "incoming")?;
match current {
None => Ok(json_number(i)),
Some(value) => {
let c = number_or_error(value, "max", "current")?;
Ok(json_number(c.max(i)))
}
}
}
fn apply_min(current: Option<&Value>, incoming: Value) -> Result<Value> {
let i = number_or_error(&incoming, "min", "incoming")?;
match current {
None => Ok(json_number(i)),
Some(value) => {
let c = number_or_error(value, "min", "current")?;
Ok(json_number(c.min(i)))
}
}
}
fn apply_merge(current: Option<&Value>, incoming: Value) -> Result<Value> {
let mut map = match current {
None => serde_json::Map::new(),
Some(Value::Object(m)) => m.clone(),
Some(other) => bail!(
"reducer 'merge' requires an object (or absent) for the current value, got {}",
type_name(other)
),
};
match incoming {
Value::Object(items) => {
for (k, v) in items {
map.insert(k, v);
}
}
other => bail!(
"reducer 'merge' requires an object for the incoming value, got {}",
type_name(&other)
),
}
Ok(Value::Object(map))
}
fn number_or_error(value: &Value, reducer_name: &str, position: &str) -> Result<f64> {
match value.as_f64() {
Some(n) => Ok(n),
None => bail!(
"reducer '{reducer_name}' requires a number for the {position} value, got {}",
type_name(value)
),
}
}
// Numeric reducers compute in f64 for simplicity. Integer typing is preserved when the result is losslessly
// representable as i64.
fn json_number(n: f64) -> Value {
if n.fract() == 0.0 && n.is_finite() && n.abs() <= (i64::MAX as f64) {
Value::Number(Number::from(n as i64))
} else {
match Number::from_f64(n) {
Some(num) => Value::Number(num),
None => Value::Null,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn append_to_absent_creates_single_element_array() {
let result = apply(Reducer::Append, None, json!("a")).unwrap();
assert_eq!(result, json!(["a"]));
}
#[test]
fn append_pushes_onto_existing_array() {
let current = json!(["a", "b"]);
let result = apply(Reducer::Append, Some(&current), json!("c")).unwrap();
assert_eq!(result, json!(["a", "b", "c"]));
}
#[test]
fn append_errors_when_current_is_not_array() {
let current = json!("not an array");
let err = apply(Reducer::Append, Some(&current), json!("x"))
.unwrap_err()
.to_string();
assert!(err.contains("'append'"), "got: {err}");
assert!(err.contains("string"), "got: {err}");
}
#[test]
fn extend_concatenates_arrays() {
let current = json!([1, 2]);
let result = apply(Reducer::Extend, Some(&current), json!([3, 4])).unwrap();
assert_eq!(result, json!([1, 2, 3, 4]));
}
#[test]
fn extend_from_absent_with_array() {
let result = apply(Reducer::Extend, None, json!([1, 2])).unwrap();
assert_eq!(result, json!([1, 2]));
}
#[test]
fn extend_errors_when_incoming_is_not_array() {
let err = apply(Reducer::Extend, None, json!(42))
.unwrap_err()
.to_string();
assert!(err.contains("'extend'"), "got: {err}");
assert!(err.contains("number"), "got: {err}");
assert!(err.contains("incoming"), "got: {err}");
}
#[test]
fn concat_joins_strings_with_newline() {
let current = json!("first");
let result = apply(Reducer::Concat, Some(&current), json!("second")).unwrap();
assert_eq!(result, json!("first\nsecond"));
}
#[test]
fn concat_from_absent_yields_incoming() {
let result = apply(Reducer::Concat, None, json!("hello")).unwrap();
assert_eq!(result, json!("hello"));
}
#[test]
fn concat_skips_separator_when_current_is_empty_string() {
let current = json!("");
let result = apply(Reducer::Concat, Some(&current), json!("first")).unwrap();
assert_eq!(result, json!("first"));
}
#[test]
fn concat_errors_when_incoming_is_not_string() {
let err = apply(Reducer::Concat, None, json!(42))
.unwrap_err()
.to_string();
assert!(err.contains("'concat'"), "got: {err}");
assert!(err.contains("number"), "got: {err}");
}
#[test]
fn sum_adds_numbers() {
let current = json!(5);
let result = apply(Reducer::Sum, Some(&current), json!(7)).unwrap();
assert_eq!(result, json!(12));
}
#[test]
fn sum_starts_from_zero_when_current_absent() {
let result = apply(Reducer::Sum, None, json!(3.5)).unwrap();
assert_eq!(result, json!(3.5));
}
#[test]
fn sum_preserves_integer_type_for_whole_results() {
let current = json!(2);
let result = apply(Reducer::Sum, Some(&current), json!(3)).unwrap();
assert!(result.is_i64(), "expected integer, got {result:?}");
assert_eq!(result, json!(5));
}
#[test]
fn sum_uses_float_when_result_has_fractional() {
let current = json!(1.5);
let result = apply(Reducer::Sum, Some(&current), json!(2.25)).unwrap();
assert_eq!(result, json!(3.75));
}
#[test]
fn sum_errors_on_string_incoming() {
let err = apply(Reducer::Sum, None, json!("not a number"))
.unwrap_err()
.to_string();
assert!(err.contains("'sum'"), "got: {err}");
assert!(err.contains("string"), "got: {err}");
}
#[test]
fn max_returns_larger_of_two() {
let current = json!(5);
let result = apply(Reducer::Max, Some(&current), json!(3)).unwrap();
assert_eq!(result, json!(5));
let result = apply(Reducer::Max, Some(&current), json!(10)).unwrap();
assert_eq!(result, json!(10));
}
#[test]
fn max_yields_incoming_when_current_absent() {
let result = apply(Reducer::Max, None, json!(42)).unwrap();
assert_eq!(result, json!(42));
}
#[test]
fn min_returns_smaller_of_two() {
let current = json!(5);
let result = apply(Reducer::Min, Some(&current), json!(3)).unwrap();
assert_eq!(result, json!(3));
let result = apply(Reducer::Min, Some(&current), json!(10)).unwrap();
assert_eq!(result, json!(5));
}
#[test]
fn min_errors_on_non_numeric_current() {
let current = json!("oops");
let err = apply(Reducer::Min, Some(&current), json!(1))
.unwrap_err()
.to_string();
assert!(err.contains("'min'"), "got: {err}");
assert!(err.contains("current"), "got: {err}");
}
#[test]
fn merge_unions_objects_with_incoming_winning_collisions() {
let current = json!({ "a": 1, "b": 2 });
let incoming = json!({ "b": 99, "c": 3 });
let result = apply(Reducer::Merge, Some(&current), incoming).unwrap();
assert_eq!(result, json!({ "a": 1, "b": 99, "c": 3 }));
}
#[test]
fn merge_from_absent_yields_incoming_object() {
let result = apply(Reducer::Merge, None, json!({ "k": "v" })).unwrap();
assert_eq!(result, json!({ "k": "v" }));
}
#[test]
fn merge_errors_when_incoming_is_not_object() {
let err = apply(Reducer::Merge, None, json!([1, 2]))
.unwrap_err()
.to_string();
assert!(err.contains("'merge'"), "got: {err}");
assert!(err.contains("array"), "got: {err}");
}
#[test]
fn merge_errors_when_current_is_not_object() {
let current = json!("not object");
let err = apply(Reducer::Merge, Some(&current), json!({ "k": "v" }))
.unwrap_err()
.to_string();
assert!(err.contains("'merge'"), "got: {err}");
assert!(err.contains("current"), "got: {err}");
}
#[test]
fn overwrite_ignores_current_and_returns_incoming() {
let current = json!("old");
let result = apply(Reducer::Overwrite, Some(&current), json!("new")).unwrap();
assert_eq!(result, json!("new"));
}
#[test]
fn overwrite_works_with_absent_current() {
let result = apply(Reducer::Overwrite, None, json!(42)).unwrap();
assert_eq!(result, json!(42));
}
}
+557
View File
@@ -0,0 +1,557 @@
use super::state::{StateManager, StateRepresentation};
use super::types::ScriptNode;
use crate::config::paths;
use crate::function::Language;
use anyhow::{Context, Result, anyhow, bail};
use serde_json::Value;
use std::collections::HashMap;
use std::env;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;
use tokio::process::Command;
use tokio::time::timeout;
#[cfg(windows)]
const PATH_SEP: &str = ";";
#[cfg(not(windows))]
const PATH_SEP: &str = ":";
#[derive(Clone)]
pub struct ScriptExecutor {
base_dir: PathBuf,
extra_envs: HashMap<String, String>,
}
impl ScriptExecutor {
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
let base_dir = base_dir.into();
let extra_envs = build_default_envs(&base_dir);
Self {
base_dir,
extra_envs,
}
}
pub fn with_envs(mut self, envs: HashMap<String, String>) -> Self {
self.extra_envs.extend(envs);
self
}
pub async fn execute(
&self,
node: &ScriptNode,
state_manager: &mut StateManager,
) -> Result<Option<String>> {
let script_path = self.base_dir.join(&node.script);
if !script_path.exists() {
bail!("Script file not found: '{}'", script_path.display());
}
let language = detect_language(&script_path)?;
let state_repr = state_manager.serialize_state()?;
let mut cmd = build_command(language, &script_path)?;
cmd.stdout(Stdio::piped());
cmd.stderr(Stdio::piped());
cmd.envs(&self.extra_envs);
match &state_repr {
StateRepresentation::Inline(json) => {
cmd.env("GRAPH_STATE", json);
}
StateRepresentation::File(path) => {
cmd.env("GRAPH_STATE_FILE", path);
}
}
let timeout_dur = Duration::from_secs(node.timeout);
let output = timeout(timeout_dur, cmd.output())
.await
.with_context(|| {
format!(
"Script '{}' timed out after {}s",
script_path.display(),
node.timeout
)
})?
.with_context(|| {
format!(
"Failed to spawn script process for '{}'",
script_path.display()
)
})?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
bail!(
"Script '{}' failed with exit code {:?}:\n{}",
script_path.display(),
output.status.code(),
stderr.trim()
);
}
let stdout = String::from_utf8_lossy(&output.stdout);
let json_output = stdout.trim();
if json_output.is_empty() {
bail!(
"Script '{}' produced no output (scripts must emit a single JSON object on stdout)",
script_path.display()
);
}
let next = state_manager
.merge_script_output(json_output)
.with_context(|| {
format!(
"Failed to merge output from script '{}'",
script_path.display()
)
})?;
apply_state_updates(node, state_manager);
Ok(next)
}
}
fn apply_state_updates(node: &ScriptNode, state_manager: &mut StateManager) {
let Some(updates) = &node.state_updates else {
return;
};
for (key, template) in updates {
let value = state_manager.interpolate_lenient(template);
state_manager
.state_mut()
.set(key.clone(), Value::String(value));
}
}
fn build_default_envs(agent_data_dir: &Path) -> HashMap<String, String> {
let mut envs = HashMap::new();
envs.insert(
"LLM_ROOT_DIR".to_string(),
paths::config_dir().to_string_lossy().into_owned(),
);
envs.insert(
"LLM_PROMPT_UTILS_FILE".to_string(),
paths::bash_prompt_utils_file()
.to_string_lossy()
.into_owned(),
);
envs.insert(
"LLM_AGENT_DATA_DIR".to_string(),
agent_data_dir.to_string_lossy().into_owned(),
);
envs.insert("CLICOLOR_FORCE".to_string(), "1".to_string());
envs.insert("FORCE_COLOR".to_string(), "1".to_string());
if let Ok(current_path) = env::var("PATH") {
let bin_dir = paths::functions_bin_dir();
envs.insert(
"PATH".to_string(),
format!("{}{}{}", bin_dir.display(), PATH_SEP, current_path),
);
}
envs
}
fn detect_language(script_path: &Path) -> Result<Language> {
let ext = script_path
.extension()
.and_then(|e| e.to_str())
.ok_or_else(|| anyhow!("Script has no file extension: '{}'", script_path.display()))?
.to_string();
match Language::from(&ext) {
Language::Unsupported => bail!(
"Unsupported script extension '.{}' for '{}'",
ext,
script_path.display()
),
lang => Ok(lang),
}
}
fn build_command(language: Language, script_path: &Path) -> Result<Command> {
let (program, prefix_args) = language.direct_invoker().ok_or_else(|| {
anyhow!(
"No direct invoker available for script '{}'",
script_path.display()
)
})?;
let mut cmd = Command::new(program);
for arg in prefix_args {
cmd.arg(arg);
}
cmd.arg(script_path);
Ok(cmd)
}
#[cfg(test)]
mod tests {
use super::super::MAX_STATE_SIZE_BYTES;
use super::*;
use crate::utils::temp_file;
use indoc::formatdoc;
use serde_json::json;
use std::collections::HashMap;
use std::env::temp_dir;
use std::fs;
fn cmd_available(name: &str) -> bool {
which::which(name).is_ok()
}
fn write_script(contents: &str, ext: &str) -> (PathBuf, PathBuf) {
let dir = temp_file("-graph-script-test-", "");
fs::create_dir_all(&dir).unwrap();
let path = dir.join(format!("script.{ext}"));
fs::write(&path, contents).unwrap();
(dir, path)
}
fn cleanup(dir: &Path) {
let _ = fs::remove_dir_all(dir);
}
fn node_for(script_filename: &str, timeout: u64) -> ScriptNode {
ScriptNode {
script: script_filename.into(),
state_updates: None,
fallback: None,
timeout,
}
}
#[tokio::test]
async fn bash_script_merges_json_output_into_state() {
if !cmd_available("bash") {
eprintln!("skipping: bash not available");
return;
}
let (dir, path) = write_script(
r#"#!/bin/bash
echo '{"quality": 0.85, "issues": 3, "_next": "approve"}'
"#,
"sh",
);
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
let next = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap();
assert_eq!(next.as_deref(), Some("approve"));
assert_eq!(state.state().get("quality"), Some(&json!(0.85)));
assert_eq!(state.state().get("issues"), Some(&json!(3)));
assert!(state.state().get("_next").is_none());
cleanup(&dir);
}
#[tokio::test]
async fn bash_script_can_read_state_from_env() {
if !cmd_available("bash") || !cmd_available("python3") {
eprintln!("skipping: bash or python3 not available");
return;
}
let (dir, path) = write_script(
r#"#!/bin/bash
NAME=$(python3 -c 'import json,os; print(json.loads(os.environ["GRAPH_STATE"])["name"])')
printf '{"greeting": "hello %s"}' "$NAME"
"#,
"sh",
);
let mut initial = HashMap::new();
initial.insert("name".into(), json!("alice"));
let mut state = StateManager::new(initial);
let executor = ScriptExecutor::new(&dir);
let _ = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap();
assert_eq!(state.state().get("greeting"), Some(&json!("hello alice")));
cleanup(&dir);
}
#[tokio::test]
async fn script_without_next_returns_none() {
if !cmd_available("bash") {
return;
}
let (dir, path) = write_script(
r#"#!/bin/bash
echo '{"ok": true}'
"#,
"sh",
);
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
let next = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap();
assert!(next.is_none());
assert_eq!(state.state().get("ok"), Some(&json!(true)));
cleanup(&dir);
}
#[tokio::test]
async fn state_updates_apply_after_json_merge() {
if !cmd_available("bash") {
return;
}
let (dir, path) = write_script(
r#"#!/bin/bash
echo '{"raw": "hello"}'
"#,
"sh",
);
let mut node = node_for(path.file_name().unwrap().to_str().unwrap(), 5);
let mut updates = HashMap::new();
updates.insert("decorated".into(), "[{{raw}}]".into());
node.state_updates = Some(updates);
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
executor.execute(&node, &mut state).await.unwrap();
assert_eq!(state.state().get("raw"), Some(&json!("hello")));
assert_eq!(state.state().get("decorated"), Some(&json!("[hello]")));
cleanup(&dir);
}
#[tokio::test]
async fn missing_script_file_errors_before_spawning() {
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(temp_dir());
let err = executor
.execute(&node_for("__does_not_exist__.sh", 5), &mut state)
.await
.unwrap_err()
.to_string();
assert!(err.contains("Script file not found"), "got: {err}");
}
#[tokio::test]
async fn empty_stdout_errors() {
if !cmd_available("bash") {
return;
}
let (dir, path) = write_script("#!/bin/bash\n", "sh");
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
let err = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap_err()
.to_string();
assert!(err.contains("produced no output"), "got: {err}");
cleanup(&dir);
}
#[tokio::test]
async fn non_json_output_errors() {
if !cmd_available("bash") {
return;
}
let (dir, path) = write_script(
&formatdoc! {r#"
#!/bin/bash
echo "not json at all"
"#},
"sh",
);
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
let err = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap_err()
.to_string();
assert!(err.contains("merge output"), "got: {err}");
cleanup(&dir);
}
#[tokio::test]
async fn non_zero_exit_errors_and_includes_stderr() {
if !cmd_available("bash") {
return;
}
let (dir, path) = write_script(
&formatdoc! {r#"
#!/bin/bash
echo "bad happened" >&2
exit 7
"#},
"sh",
);
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
let err = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap_err()
.to_string();
assert!(err.contains("exit code"), "got: {err}");
assert!(err.contains("bad happened"), "got: {err}");
cleanup(&dir);
}
#[tokio::test]
async fn execution_timeout_is_enforced() {
if !cmd_available("bash") {
return;
}
let (dir, path) = write_script(
r#"#!/bin/bash
sleep 5
echo '{"ok":true}'
"#,
"sh",
);
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
let err = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 1),
&mut state,
)
.await
.unwrap_err()
.to_string();
assert!(err.contains("timed out"), "got: {err}");
cleanup(&dir);
}
#[tokio::test]
async fn large_state_is_delivered_via_file_env_var() {
if !cmd_available("bash") || !cmd_available("python3") {
return;
}
let big = "x".repeat(MAX_STATE_SIZE_BYTES + 1024);
let mut initial = HashMap::new();
initial.insert("blob".into(), json!(big));
let (dir, path) = write_script(
r#"#!/bin/bash
if [ -n "$GRAPH_STATE_FILE" ]; then
LEN=$(python3 -c 'import json,os; print(len(json.load(open(os.environ["GRAPH_STATE_FILE"]))["blob"]))')
printf '{"blob_len": %s, "via_file": true}' "$LEN"
elif [ -n "$GRAPH_STATE" ]; then
echo '{"via_file": false}'
fi
"#,
"sh",
);
let mut state = StateManager::new(initial);
let executor = ScriptExecutor::new(&dir);
executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 10),
&mut state,
)
.await
.unwrap();
assert_eq!(state.state().get("via_file"), Some(&json!(true)));
let len = state.state().get("blob_len").unwrap().as_i64().unwrap();
assert_eq!(len as usize, big.len());
cleanup(&dir);
}
#[tokio::test]
async fn python_script_can_emit_routing_and_state() {
if !cmd_available("python3") {
eprintln!("skipping: python3 not available");
return;
}
let (dir, path) = write_script(
r#"import os, json
state = json.loads(os.environ["GRAPH_STATE"])
print(json.dumps({
"_next": "next_node",
"doubled": state.get("n", 0) * 2,
}))
"#,
"py",
);
let mut initial = HashMap::new();
initial.insert("n".into(), json!(21));
let mut state = StateManager::new(initial);
let executor = ScriptExecutor::new(&dir);
let next = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap();
assert_eq!(next.as_deref(), Some("next_node"));
assert_eq!(state.state().get("doubled"), Some(&json!(42)));
cleanup(&dir);
}
#[tokio::test]
async fn unknown_extension_is_rejected() {
let (dir, path) = write_script("echo hi", "xyz");
let mut state = StateManager::new(HashMap::new());
let executor = ScriptExecutor::new(&dir);
let err = executor
.execute(
&node_for(path.file_name().unwrap().to_str().unwrap(), 5),
&mut state,
)
.await
.unwrap_err()
.to_string();
assert!(
err.contains("Unsupported script extension '.xyz'"),
"got: {err}"
);
cleanup(&dir);
}
}
+9
View File
@@ -0,0 +1,9 @@
use serde_json::Value;
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct BranchWrites {
pub node_id: String,
pub invocation_index: usize,
pub writes: HashMap<String, Value>,
}
+1021
View File
File diff suppressed because it is too large Load Diff
+189
View File
@@ -0,0 +1,189 @@
use crate::client::call_chat_completions;
use crate::config::{Input, RequestContext, Role, RoleLike};
use crate::utils::create_abort_signal;
use anyhow::{Context, Result, bail};
use serde_json::Value;
use std::sync::Arc;
const EXTRACTOR_ROLE_NAME: &str = "__structured_output__";
const EXTRACTOR_ROLE_PROMPT: &str = "\
Extract a JSON object from the user's input that strictly conforms to the provided JSON Schema.
Rules:
- Output ONLY the JSON object. No prose, no explanation, no markdown fences, no <think> tokens.
- The first character of your response must be `{` and the last must be `}`.
- Every key marked `required` in the schema MUST appear in the output.
- All values MUST match the types specified in the schema.
- If the input is already a valid JSON object matching the schema, return it unchanged.
- If a field cannot be determined from the input, use `null` (when allowed) or your best inferred value.
- Do NOT invent fields not present in the schema.";
pub async fn extract(raw: &str, schema: &Value, parent_ctx: &mut RequestContext) -> Result<Value> {
if let Some(parsed) = try_parse_json(raw) {
return Ok(parsed);
}
extract_via_extractor(raw, schema, parent_ctx, false).await
}
async fn extract_via_extractor(
raw: &str,
schema: &Value,
parent_ctx: &mut RequestContext,
is_repair: bool,
) -> Result<Value> {
let role = build_extractor_role()?;
let prompt = build_extractor_prompt(raw, schema, is_repair);
let saved_role = parent_ctx.role.clone();
parent_ctx.role = Some(role);
let result = run_one_shot(&prompt, parent_ctx).await;
parent_ctx.role = saved_role;
let output = result.context("Structured-output extractor LLM call failed")?;
match try_parse_json(&output) {
Some(value) => Ok(value),
None if is_repair => bail!(
"Structured-output extractor failed to produce valid JSON after repair retry. \
Last response:\n{output}"
),
None => Box::pin(extract_via_extractor(&output, schema, parent_ctx, true)).await,
}
}
fn build_extractor_role() -> Result<Role> {
let mut role = Role::new(EXTRACTOR_ROLE_NAME, EXTRACTOR_ROLE_PROMPT);
role.set_enabled_tools(Some(String::new()));
role.set_enabled_mcp_servers(Some(String::new()));
Ok(role)
}
fn build_extractor_prompt(raw: &str, schema: &Value, is_repair: bool) -> String {
let schema_json = serde_json::to_string_pretty(schema).unwrap_or_else(|_| schema.to_string());
if is_repair {
format!(
"Your previous response was not valid JSON. Output ONLY a JSON object \
matching this schema. No prose, no fences.\n\nSchema:\n{schema_json}\n\nInput:\n{raw}"
)
} else {
format!("Schema:\n{schema_json}\n\nInput:\n{raw}")
}
}
async fn run_one_shot(prompt: &str, ctx: &mut RequestContext) -> Result<String> {
let abort = create_abort_signal();
let app_cfg = Arc::clone(&ctx.app.config);
let role_for_input = ctx.role.clone();
let input = Input::from_str(ctx, prompt, role_for_input);
let client = input.create_client()?;
ctx.before_chat_completion(&input)?;
let (output, tool_results) =
call_chat_completions(&input, false, false, client.as_ref(), ctx, abort).await?;
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
Ok(output)
}
fn try_parse_json(raw: &str) -> Option<Value> {
let cleaned = strip_code_fences(raw.trim());
serde_json::from_str(cleaned).ok()
}
fn strip_code_fences(s: &str) -> &str {
let after_open = s
.strip_prefix("```json")
.or_else(|| s.strip_prefix("```"))
.map(str::trim_start)
.unwrap_or(s);
after_open
.strip_suffix("```")
.map(str::trim_end)
.unwrap_or(after_open)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn try_parse_json_accepts_plain_object() {
let v = try_parse_json(r#"{"a": 1}"#).unwrap();
assert_eq!(v, json!({"a": 1}));
}
#[test]
fn try_parse_json_strips_json_fences() {
let raw = "```json\n{\"a\": 1}\n```";
let v = try_parse_json(raw).unwrap();
assert_eq!(v, json!({"a": 1}));
}
#[test]
fn try_parse_json_strips_bare_fences() {
let raw = "```\n{\"a\": 1}\n```";
let v = try_parse_json(raw).unwrap();
assert_eq!(v, json!({"a": 1}));
}
#[test]
fn try_parse_json_tolerates_whitespace() {
let v = try_parse_json(" \n {\"x\": true}\n\n").unwrap();
assert_eq!(v, json!({"x": true}));
}
#[test]
fn try_parse_json_returns_none_on_prose() {
assert!(try_parse_json("Here is the result: it's good").is_none());
}
#[test]
fn try_parse_json_returns_none_on_partial_json() {
assert!(try_parse_json("{\"a\": ").is_none());
}
#[test]
fn try_parse_json_accepts_arrays() {
let v = try_parse_json("[1, 2, 3]").unwrap();
assert_eq!(v, json!([1, 2, 3]));
}
#[test]
fn build_extractor_prompt_includes_schema_and_input() {
let schema = json!({"type": "object"});
let prompt = build_extractor_prompt("hello", &schema, false);
assert!(prompt.contains("Schema:"));
assert!(prompt.contains("Input:"));
assert!(prompt.contains("hello"));
}
#[test]
fn build_extractor_prompt_repair_includes_repair_instruction() {
let schema = json!({"type": "object"});
let prompt = build_extractor_prompt("oops", &schema, true);
assert!(prompt.contains("previous response"));
assert!(prompt.contains("oops"));
}
#[test]
fn build_extractor_role_disables_tools_and_mcp() {
let role = build_extractor_role().expect("builtin role must exist");
assert_eq!(role.enabled_tools().as_deref(), Some(""));
assert_eq!(role.enabled_mcp_servers().as_deref(), Some(""));
}
}
+1116
View File
File diff suppressed because it is too large Load Diff
+369
View File
@@ -0,0 +1,369 @@
use super::state::StateManager;
use super::types::{ApprovalNode, InputNode};
use crate::config::RequestContext;
use crate::function::user_interaction::{USER_FUNCTION_PREFIX, handle_user_tool};
use anyhow::{Context, Result, anyhow, bail};
use serde_json::{Value, json};
use std::collections::HashMap;
const CHOICE_KEY: &str = "choice";
const INPUT_KEY: &str = "input";
pub struct ApprovalNodeExecutor;
impl ApprovalNodeExecutor {
pub async fn execute(
node: &ApprovalNode,
state_manager: &mut StateManager,
ctx: &mut RequestContext,
) -> Result<String> {
let question = state_manager
.interpolate(&node.question)
.context("Failed to interpolate approval question")?;
let response = handle_user_tool(
ctx,
&format!("{USER_FUNCTION_PREFIX}ask"),
&json!({ "question": question, "options": node.options }),
)
.await
.context("user__ask failed")?;
if let Some(err) = response.get("error").and_then(Value::as_str) {
bail!("Approval interaction failed: {err}");
}
let choice = response
.get("answer")
.and_then(Value::as_str)
.context("Approval response missing 'answer' field")?
.to_string();
apply_state_updates_with_var(&node.state_updates, state_manager, CHOICE_KEY, &choice);
resolve_approval_route(node, &choice)
}
}
pub struct InputNodeExecutor;
impl InputNodeExecutor {
pub async fn execute(
node: &InputNode,
node_next: Option<&str>,
state_manager: &mut StateManager,
ctx: &mut RequestContext,
) -> Result<String> {
let question = build_input_question(node, state_manager)?;
let response = handle_user_tool(
ctx,
&format!("{USER_FUNCTION_PREFIX}input"),
&json!({ "question": question }),
)
.await
.context("user__input failed")?;
if let Some(err) = response.get("error").and_then(Value::as_str) {
bail!("Input interaction failed: {err}");
}
let raw = response
.get("answer")
.and_then(Value::as_str)
.context("Input response missing 'answer' field")?
.to_string();
let input_text = if raw.is_empty() {
node.default
.as_ref()
.map(|t| state_manager.interpolate_lenient(t))
.unwrap_or_default()
} else {
raw
};
if let Some(expr) = &node.validation
&& !validate_length(&input_text, expr)?
{
bail!(
"Input failed validation '{}' (got {} chars)",
expr,
input_text.chars().count()
);
}
apply_state_updates_with_var(&node.state_updates, state_manager, INPUT_KEY, &input_text);
node_next
.map(String::from)
.ok_or_else(|| anyhow!("Input node has no `next` set"))
}
}
fn build_input_question(node: &InputNode, state_manager: &StateManager) -> Result<String> {
let mut question = state_manager
.interpolate(&node.question)
.context("Failed to interpolate input question")?;
if let Some(default_template) = &node.default {
let default = state_manager.interpolate_lenient(default_template);
if !default.is_empty() {
question = format!("{question} [default: {default}]");
}
}
Ok(question)
}
fn resolve_approval_route(node: &ApprovalNode, choice: &str) -> Result<String> {
if let Some(target) = node.routes.get(choice) {
return Ok(target.clone());
}
Ok(node.on_other.clone())
}
fn apply_state_updates_with_var(
updates: &Option<HashMap<String, String>>,
state_manager: &mut StateManager,
var_name: &str,
var_value: &str,
) {
let Some(updates) = updates else {
return;
};
let prev = state_manager.state().get(var_name).cloned();
state_manager
.state_mut()
.set(var_name.into(), Value::String(var_value.to_string()));
for (key, template) in updates {
let value = state_manager.interpolate_lenient(template);
state_manager
.state_mut()
.set(key.clone(), Value::String(value));
}
match prev {
Some(v) => state_manager.state_mut().set(var_name.into(), v),
None => {
state_manager.state_mut().set(var_name.into(), Value::Null);
}
}
}
/// Evaluate a `len(input) OP N` expression where OP is one of `>`, `>=`,
/// `<`, `<=`, `==`. Lengths are byte counts (matches Rust's `str::len`).
/// Other expressions are rejected at runtime.
fn validate_length(input: &str, expr: &str) -> Result<bool> {
let trimmed = expr.trim();
let after_len = trimmed
.strip_prefix("len(input)")
.map(str::trim)
.ok_or_else(|| {
anyhow!(
"Unsupported validation expression '{expr}'; only `len(input) OP N` is supported"
)
})?;
let (op, rhs_str) = if let Some(rest) = after_len.strip_prefix(">=") {
(">=", rest)
} else if let Some(rest) = after_len.strip_prefix("<=") {
("<=", rest)
} else if let Some(rest) = after_len.strip_prefix("==") {
("==", rest)
} else if let Some(rest) = after_len.strip_prefix('>') {
(">", rest)
} else if let Some(rest) = after_len.strip_prefix('<') {
("<", rest)
} else {
bail!("No comparison operator in validation expression '{expr}'");
};
let rhs: usize = rhs_str
.trim()
.parse()
.with_context(|| format!("Invalid right-hand side in validation '{expr}'"))?;
let len = input.len();
Ok(match op {
">=" => len >= rhs,
"<=" => len <= rhs,
"==" => len == rhs,
">" => len > rhs,
"<" => len < rhs,
_ => unreachable!(),
})
}
#[cfg(test)]
mod tests {
use super::super::types::*;
use super::*;
use serde_json::json;
fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
let mut map = HashMap::new();
for (k, v) in pairs {
map.insert((*k).into(), v.clone());
}
StateManager::new(map)
}
fn approval(options: &[&str], routes: &[(&str, &str)], on_other: &str) -> ApprovalNode {
let mut r = HashMap::new();
for (k, v) in routes {
r.insert((*k).into(), (*v).into());
}
ApprovalNode {
question: "?".into(),
options: options.iter().map(|s| (*s).into()).collect(),
routes: r,
on_other: on_other.into(),
state_updates: None,
}
}
fn input(question: &str) -> InputNode {
InputNode {
question: question.into(),
default: None,
validation: None,
state_updates: None,
}
}
#[test]
fn validate_length_supports_all_comparison_operators() {
assert!(validate_length("hello", "len(input) > 0").unwrap());
assert!(!validate_length("", "len(input) > 0").unwrap());
assert!(validate_length("hello", "len(input) >= 5").unwrap());
assert!(!validate_length("hi", "len(input) >= 5").unwrap());
assert!(validate_length("hello", "len(input) < 10").unwrap());
assert!(!validate_length("hello world!", "len(input) < 10").unwrap());
assert!(validate_length("hi", "len(input) <= 2").unwrap());
assert!(!validate_length("hello", "len(input) <= 2").unwrap());
assert!(validate_length("hello", "len(input) == 5").unwrap());
assert!(!validate_length("hello", "len(input) == 3").unwrap());
}
#[test]
fn validate_length_handles_whitespace() {
assert!(validate_length("hi", " len(input) >= 1 ").unwrap());
}
#[test]
fn validate_length_rejects_unsupported_expressions() {
assert!(validate_length("x", "matches /[a-z]+/").is_err());
assert!(validate_length("x", "len(input)").is_err());
assert!(validate_length("x", "len(input) >").is_err());
assert!(validate_length("x", "len(input) >= abc").is_err());
}
#[test]
fn approval_route_lookup_returns_target_on_match() {
let node = approval(
&["yes", "no"],
&[("yes", "deploy"), ("no", "cancel")],
"clarify",
);
assert_eq!(resolve_approval_route(&node, "yes").unwrap(), "deploy");
assert_eq!(resolve_approval_route(&node, "no").unwrap(), "cancel");
}
#[test]
fn approval_route_lookup_falls_back_to_on_other_for_unknown_choice() {
let node = approval(
&["yes", "no"],
&[("yes", "deploy"), ("no", "cancel")],
"clarify",
);
assert_eq!(resolve_approval_route(&node, "maybe").unwrap(), "clarify");
assert_eq!(
resolve_approval_route(&node, "free-form text").unwrap(),
"clarify"
);
}
#[test]
fn state_updates_expose_choice_during_evaluation_only() {
let mut updates = HashMap::new();
updates.insert("decision".into(), "{{choice}}".into());
let mut state = manager_with(&[]);
apply_state_updates_with_var(&Some(updates), &mut state, CHOICE_KEY, "approve");
assert_eq!(state.state().get("decision"), Some(&json!("approve")));
assert_eq!(state.state().get(CHOICE_KEY), Some(&Value::Null));
}
#[test]
fn state_updates_preserve_pre_existing_var_value() {
let mut updates = HashMap::new();
updates.insert("decision".into(), "{{choice}}".into());
let mut state = manager_with(&[("choice", json!("preserved"))]);
apply_state_updates_with_var(&Some(updates), &mut state, CHOICE_KEY, "approve");
assert_eq!(state.state().get("decision"), Some(&json!("approve")));
assert_eq!(state.state().get(CHOICE_KEY), Some(&json!("preserved")));
}
#[test]
fn state_updates_for_input_use_input_key() {
let mut updates = HashMap::new();
updates.insert("api_key".into(), "{{input}}".into());
let mut state = manager_with(&[]);
apply_state_updates_with_var(&Some(updates), &mut state, INPUT_KEY, "sk-12345");
assert_eq!(state.state().get("api_key"), Some(&json!("sk-12345")));
assert_eq!(state.state().get(INPUT_KEY), Some(&Value::Null));
}
#[test]
fn input_question_appends_default_when_present() {
let state = manager_with(&[("name", json!("alice"))]);
let mut node = input("Hi, what's your name?");
node.default = Some("{{name}}".into());
let q = build_input_question(&node, &state).unwrap();
assert_eq!(q, "Hi, what's your name? [default: alice]");
}
#[test]
fn input_question_omits_default_when_blank_after_interpolation() {
let state = manager_with(&[]);
let mut node = input("Enter value:");
node.default = Some("{{missing}}".into());
let q = build_input_question(&node, &state).unwrap();
assert_eq!(q, "Enter value:");
}
#[test]
fn input_question_uses_no_default_when_field_absent() {
let state = manager_with(&[]);
let node = input("Enter value:");
let q = build_input_question(&node, &state).unwrap();
assert_eq!(q, "Enter value:");
}
#[test]
fn no_state_updates_means_var_never_appears_in_state() {
let mut state = manager_with(&[]);
apply_state_updates_with_var(&None, &mut state, CHOICE_KEY, "approve");
assert!(state.state().get(CHOICE_KEY).is_none());
}
}
File diff suppressed because it is too large Load Diff
+27
View File
@@ -2,6 +2,7 @@ mod cli;
mod client;
mod config;
mod function;
mod graph;
mod rag;
mod render;
mod repl;
@@ -82,8 +83,23 @@ async fn main() -> Result<()> {
let log_path = setup_logger()?;
if let Some(version) = &cli.update {
let version = version.clone();
let force = cli.force;
return tokio::task::spawn_blocking(move || config::run_self_update(version, force))
.await?;
}
install_builtins()?;
if let Some(category) = cli.install {
return config::install_assets(category);
}
if let Some(url) = cli.install_from.as_deref() {
return config::install_remote(url, cli.filter, cli.install_force);
}
if let Some(client_arg) = &cli.authenticate {
let cfg = Config::load_with_interpolation(true).await?;
let app_config = AppConfig::from_config(cfg)?;
@@ -311,6 +327,17 @@ async fn start_directive(
abort_signal: AbortSignal,
) -> Result<()> {
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
if graph::active_agent_graph_name(ctx).is_some() {
ctx.before_chat_completion(&input)?;
let output =
graph::run_active_agent_graph(ctx, &input.text(), abort_signal.clone()).await?;
app.print_markdown(&output)?;
ctx.after_chat_completion(app.as_ref(), &input, &output, &[])?;
ctx.exit_session()?;
return Ok(());
}
let client = input.create_client()?;
let extract_code = !*IS_STDOUT_TERMINAL && code_mode;
ctx.before_chat_completion(&input)?;
+18 -11
View File
@@ -8,6 +8,7 @@ use crate::vault::interpolate_secrets;
use anyhow::{Context, Result, anyhow};
use futures_util::{StreamExt, TryStreamExt, stream};
use http::{HeaderName, HeaderValue};
use indexmap::IndexMap;
use indoc::formatdoc;
use rmcp::service::RunningService;
use rmcp::transport::StreamableHttpClientTransport;
@@ -49,23 +50,29 @@ impl Clone for ServerCatalog {
}
}
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Deserialize, Serialize)]
pub(crate) struct McpServersConfig {
#[serde(rename = "mcpServers")]
pub mcp_servers: HashMap<String, McpServer>,
pub mcp_servers: IndexMap<String, McpServer>,
}
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub(crate) struct McpServer {
#[serde(rename = "type")]
pub transport_type: McpTransportType,
#[serde(skip_serializing_if = "Option::is_none")]
pub command: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub args: Option<Vec<String>>,
pub env: Option<HashMap<String, JsonField>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub env: Option<IndexMap<String, JsonField>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cwd: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub url: Option<String>,
pub headers: Option<HashMap<String, String>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub headers: Option<IndexMap<String, String>>,
}
impl McpServer {
@@ -111,7 +118,7 @@ impl McpServer {
}
}
#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Hash)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Hash)]
#[serde(rename_all = "lowercase")]
pub(crate) enum McpTransportType {
Stdio,
@@ -119,7 +126,7 @@ pub(crate) enum McpTransportType {
Sse,
}
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(untagged)]
pub(crate) enum JsonField {
Str(String),
@@ -352,7 +359,7 @@ pub(crate) async fn spawn_mcp_server(
async fn spawn_http_mcp_server(
url: &str,
headers: Option<&HashMap<String, String>>,
headers: Option<&IndexMap<String, String>>,
) -> Result<Arc<ConnectedServer>> {
let transport = if let Some(hdrs) = headers
&& !hdrs.is_empty()
@@ -382,7 +389,7 @@ async fn spawn_http_mcp_server(
async fn spawn_sse_mcp_server(
url: &str,
headers: Option<&HashMap<String, String>>,
headers: Option<&IndexMap<String, String>>,
) -> Result<Arc<ConnectedServer>> {
let sse = LegacySseTransport::connect(url, headers)
.await
@@ -482,7 +489,7 @@ mod tests {
}
fn make_registry_with_config(server_names: &[&str]) -> McpRegistry {
let mut mcp_servers = HashMap::new();
let mut mcp_servers = IndexMap::new();
for name in server_names {
mcp_servers.insert(name.to_string(), stdio_server("echo"));
}
@@ -530,7 +537,7 @@ mod tests {
#[test]
fn validate_stdio_with_headers_fails() {
let mut headers = HashMap::new();
let mut headers = IndexMap::new();
headers.insert("Auth".into(), "Bearer tok".into());
let spec = McpServer {
transport_type: McpTransportType::Stdio,
+24 -14
View File
@@ -1,13 +1,14 @@
use anyhow::{Context, Result, anyhow};
use eventsource_stream::{EventStream, Eventsource};
use fmt::{Display, Formatter};
use futures_util::StreamExt;
use futures_util::stream::BoxStream;
use indexmap::IndexMap;
use mpsc::error::SendError;
use mpsc::{OwnedPermit, Receiver, Sender, channel};
use reqwest::Client;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use reqwest_eventsource::{Event, EventSource};
use reqwest::{Client, header};
use rmcp::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::future::Future;
@@ -17,6 +18,8 @@ use tokio::sync::mpsc;
use tokio::time::Duration;
use url::Url;
type SseEventStream = EventStream<BoxStream<'static, reqwest::Result<bytes::Bytes>>>;
const CHANNEL_BUF: usize = 64;
pub struct LegacySseTransport {
@@ -25,7 +28,10 @@ pub struct LegacySseTransport {
}
impl LegacySseTransport {
pub async fn connect(sse_url: &str, headers: Option<&HashMap<String, String>>) -> Result<Self> {
pub async fn connect(
sse_url: &str,
headers: Option<&IndexMap<String, String>>,
) -> Result<Self> {
let base_url =
Url::parse(sse_url).with_context(|| format!("Invalid SSE URL: {sse_url}"))?;
@@ -47,8 +53,15 @@ impl LegacySseTransport {
.build()
.context("Failed to build HTTP client")?;
let request = client.get(sse_url);
let mut es = EventSource::new(request).context("Failed to open SSE connection")?;
let response = client
.get(sse_url)
.header(header::ACCEPT, "text/event-stream")
.send()
.await
.context("Failed to open SSE connection")?
.error_for_status()
.context("SSE server returned an error status")?;
let mut es: SseEventStream = response.bytes_stream().boxed().eventsource();
let post_endpoint = wait_for_endpoint_event(&mut es, &base_url).await?;
@@ -83,18 +96,17 @@ impl LegacySseTransport {
}
}
async fn wait_for_endpoint_event(es: &mut EventSource, base_url: &Url) -> Result<String> {
async fn wait_for_endpoint_event(es: &mut SseEventStream, base_url: &Url) -> Result<String> {
let timeout = Duration::from_secs(30);
tokio::time::timeout(timeout, async {
while let Some(event) = es.next().await {
match event {
Ok(Event::Open) => {}
Ok(Event::Message(msg)) if msg.event == "endpoint" => {
Ok(msg) if msg.event == "endpoint" => {
let endpoint = msg.data.trim().to_string();
let resolved = resolve_endpoint(&endpoint, base_url)?;
return Ok(resolved);
}
Ok(Event::Message(_)) => {}
Ok(_) => {}
Err(e) => {
return Err(anyhow!(
"SSE connection error while waiting for endpoint event: {e}"
@@ -120,10 +132,10 @@ fn resolve_endpoint(endpoint: &str, base_url: &Url) -> Result<String> {
}
}
async fn sse_reader_task(mut es: EventSource, tx: Sender<ServerJsonRpcMessage>) {
async fn sse_reader_task(mut es: SseEventStream, tx: Sender<ServerJsonRpcMessage>) {
while let Some(event) = es.next().await {
match event {
Ok(Event::Message(msg)) if msg.event == "message" => {
Ok(msg) if msg.event == "message" => {
match serde_json::from_str::<ServerJsonRpcMessage>(&msg.data) {
Ok(rpc_msg) => {
if tx.send(rpc_msg).await.is_err() {
@@ -136,14 +148,12 @@ async fn sse_reader_task(mut es: EventSource, tx: Sender<ServerJsonRpcMessage>)
}
}
Ok(_) => {}
Err(reqwest_eventsource::Error::StreamEnded) => break,
Err(e) => {
error!("SSE stream error: {e}");
break;
}
}
}
es.close();
}
async fn post_writer_task(
+135 -11
View File
@@ -16,7 +16,8 @@ use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::{
collections::HashMap, env, fmt::Debug, fs, hash::Hash, path::Path, sync::Arc, time::Duration,
collections::HashMap, env, fmt, fmt::Debug, fs, hash::Hash, path::Path, sync::Arc,
time::Duration,
};
use tokio::time::sleep;
@@ -56,7 +57,7 @@ pub struct Rag {
}
impl Debug for Rag {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Rag")
.field("name", &self.name)
.field("path", &self.path)
@@ -81,11 +82,126 @@ impl Clone for Rag {
}
}
#[derive(Debug, Clone, Default)]
pub struct RagInitConfig {
pub embedding_model: Option<String>,
pub chunk_size: Option<usize>,
pub chunk_overlap: Option<usize>,
pub reranker_model: Option<String>,
pub top_k: Option<usize>,
pub batch_size: Option<usize>,
}
impl Rag {
fn create_embeddings_client(&self, model: Model) -> Result<Box<dyn Client>> {
init_client(&self.app_config, model)
}
pub async fn init_with_config(
app: &AppConfig,
name: &str,
save_path: &Path,
doc_paths: &[String],
config: &RagInitConfig,
abort_signal: AbortSignal,
) -> Result<Self> {
if doc_paths.is_empty() {
bail!("Cannot build RAG knowledge base '{name}' with no documents");
}
println!("⚙ Initializing RAG...");
let data = Self::resolve_init_data(app, config)?;
let mut rag = Self::create(app, name, save_path, data)?;
let loaders = app.document_loaders.clone();
let (spinner, spinner_rx) = Spinner::create("");
abortable_run_with_spinner_rx(
rag.sync_documents(doc_paths, true, loaders, Some(spinner)),
spinner_rx,
abort_signal,
)
.await?;
if rag.save()? {
println!("✓ Saved RAG to '{}'.", save_path.display());
}
Ok(rag)
}
fn resolve_init_data(app: &AppConfig, config: &RagInitConfig) -> Result<RagData> {
let embedding_model_id = config
.embedding_model
.clone()
.or_else(|| app.rag_embedding_model.clone());
let embedding_model_id = match embedding_model_id {
Some(value) => {
println!("Embedding model: {value}");
value
}
None => {
if !*IS_STDOUT_TERMINAL {
bail!(
"RAG knowledge base needs an embedding model. Set `embedding_model` \
on the rag node, or run the agent interactively once."
);
}
let models = list_models(app, ModelType::Embedding);
if models.is_empty() {
bail!("No available embedding model");
}
select_embedding_model(&models)?
}
};
let embedding_model =
Model::retrieve_model(app, &embedding_model_id, ModelType::Embedding)?;
let chunk_size = match config.chunk_size.or(app.rag_chunk_size) {
Some(value) => {
println!("Chunk size: {value}");
value
}
None => {
if !*IS_STDOUT_TERMINAL {
bail!(
"RAG knowledge base needs a chunk_size. Set `chunk_size` on the \
rag node, or run the agent interactively once."
);
}
set_chunk_size(&embedding_model)?
}
};
let chunk_overlap = match config.chunk_overlap.or(app.rag_chunk_overlap) {
Some(value) => {
println!("Chunk overlap: {value}");
value
}
None => {
if !*IS_STDOUT_TERMINAL {
bail!(
"RAG knowledge base needs a chunk_overlap. Set `chunk_overlap` on \
the rag node, or run the agent interactively once."
);
}
set_chunk_overlay(chunk_size / 20)?
}
};
let reranker_model = config
.reranker_model
.clone()
.or_else(|| app.rag_reranker_model.clone());
let top_k = config.top_k.unwrap_or(app.rag_top_k);
let batch_size = config
.batch_size
.or_else(|| embedding_model.max_batch_size());
Ok(RagData::new(
embedding_model.id(),
chunk_size,
chunk_overlap,
reranker_model,
top_k,
batch_size,
))
}
pub async fn init(
app: &AppConfig,
name: &str,
@@ -315,6 +431,14 @@ impl Rag {
self.name == TEMP_RAG_NAME
}
pub fn configured_top_k(&self) -> usize {
self.data.top_k
}
pub fn configured_reranker(&self) -> Option<&str> {
self.data.reranker_model.as_deref()
}
pub async fn search(
&self,
text: &str,
@@ -323,7 +447,7 @@ impl Rag {
abort_signal: AbortSignal,
) -> Result<(String, String, Vec<DocumentId>)> {
let ret = abortable_run_with_spinner(
self.hybird_search(text, top_k, rerank_model),
self.hybrid_search(text, top_k, rerank_model),
"Searching",
abort_signal,
)
@@ -583,7 +707,7 @@ impl Rag {
Ok(())
}
async fn hybird_search(
async fn hybrid_search(
&self,
query: &str,
top_k: usize,
@@ -781,7 +905,7 @@ pub struct RagData {
}
impl Debug for RagData {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RagData")
.field("embedding_model", &self.embedding_model)
.field("chunk_size", &self.chunk_size)
@@ -909,7 +1033,7 @@ pub type FileId = usize;
pub struct DocumentId(usize);
impl Debug for DocumentId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let (file_index, document_index) = self.split();
f.write_fmt(format_args!("{file_index}-{document_index}"))
}
@@ -951,8 +1075,8 @@ impl SelectOption {
}
}
impl std::fmt::Display for SelectOption {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
impl fmt::Display for SelectOption {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{} ({})", self.value, self.description)
}
}
@@ -1256,13 +1380,13 @@ mod tests {
#[test]
fn get_separators_returns_language_specific() {
let rs_seps = splitter::get_separators("rs");
let rs_seps = get_separators("rs");
assert!(rs_seps.iter().any(|s| s.contains("fn ")));
let py_seps = splitter::get_separators("py");
let py_seps = get_separators("py");
assert!(py_seps.iter().any(|s| s.contains("def ")));
let md_seps = splitter::get_separators("md");
let md_seps = get_separators("md");
assert!(md_seps.iter().any(|s| s.contains("# ")));
}
+20
View File
@@ -17,7 +17,11 @@ pub async fn render_stream(
rx: UnboundedReceiver<SseEvent>,
app: &AppConfig,
abort_signal: AbortSignal,
silent: bool,
) -> Result<()> {
if silent {
return drain_silently(rx, &abort_signal).await;
}
let ret = if *IS_STDOUT_TERMINAL && app.highlight {
let render_options = app.render_options()?;
let mut render = MarkdownRender::init(render_options)?;
@@ -28,6 +32,22 @@ pub async fn render_stream(
ret.map_err(|err| err.context("Failed to reader stream"))
}
async fn drain_silently(
mut rx: UnboundedReceiver<SseEvent>,
abort_signal: &AbortSignal,
) -> Result<()> {
loop {
if abort_signal.aborted() {
break;
}
match rx.recv().await {
Some(SseEvent::Done) | None => break,
Some(SseEvent::Text(_)) => {}
}
}
Ok(())
}
pub fn render_error(err: anyhow::Error) {
eprintln!("{}", error_text(&pretty_error(&err)));
}
+115 -36
View File
@@ -7,20 +7,22 @@ use self::highlighter::ReplHighlighter;
use self::prompt::ReplPrompt;
use crate::client::{call_chat_completions, call_chat_completions_streaming, init_client, oauth};
use crate::config::paths;
use crate::config::{
AgentVariables, AppConfig, AssertState, Input, LastMessage, RequestContext, StateFlags,
macro_execute,
};
use crate::config::{AssetCategory, paths};
use crate::render::render_error;
use crate::utils::{
AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file,
};
use crate::resolve_oauth_client;
use crate::{config, graph, resolve_oauth_client};
use anyhow::{Context, Result, bail};
use crossterm::cursor::SetCursorStyle;
use fancy_regex::Regex;
use indoc::indoc;
use log::warn;
use parking_lot::RwLock;
use reedline::CursorConfig;
use reedline::{
@@ -31,10 +33,20 @@ use reedline::{
use reedline::{MenuBuilder, Signal};
use std::sync::LazyLock;
use std::{env, process, sync::Arc};
use tokio::task;
const MENU_NAME: &str = "completion_menu";
static REPL_COMMANDS: LazyLock<[ReplCommand; 39]> = LazyLock::new(|| {
pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {"
[SYSTEM REMINDER - TODO CONTINUATION]
You have incomplete tasks. Rules:
1. BEFORE marking a todo done: verify the work compiles/works. No premature completion.
2. If a todo is broad (e.g. \"implement X and implement Y\"): break it into specific subtasks FIRST using todo__add, then work on those.\n\
3. Each todo should be atomic and be \"single responsibility\" - completable in one focused action.
4. Continue with the next pending item now. Call tools immediately."
};
static REPL_COMMANDS: LazyLock<[ReplCommand; 42]> = LazyLock::new(|| {
[
ReplCommand::new(".help", "Show this help guide", AssertState::pass()),
ReplCommand::new(".info", "Show system info", AssertState::pass()),
@@ -48,6 +60,11 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 39]> = LazyLock::new(|| {
"Modify configuration file",
AssertState::False(StateFlags::AGENT),
),
ReplCommand::new(
".edit mcp-config",
"Modify the MCP servers configuration file",
AssertState::False(StateFlags::AGENT),
),
ReplCommand::new(".model", "Switch LLM model", AssertState::pass()),
ReplCommand::new(
".prompt",
@@ -141,7 +158,7 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 39]> = LazyLock::new(|| {
ReplCommand::new(
".clear todo",
"Clear the todo list and stop auto-continuation",
AssertState::True(StateFlags::AGENT),
AssertState::pass(),
),
ReplCommand::new(
".rag",
@@ -201,6 +218,16 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 39]> = LazyLock::new(|| {
"View or modify the Loki vault",
AssertState::pass(),
),
ReplCommand::new(
".install",
"Reinstall bundled assets, or install assets from a remote git repo (.install remote <url>)",
AssertState::pass(),
),
ReplCommand::new(
".update",
"Update Loki to the latest release (or a specified version)",
AssertState::pass(),
),
ReplCommand::new(".exit", "Exit REPL", AssertState::pass()),
]
});
@@ -487,6 +514,12 @@ pub async fn run_repl_command(
),
},
".session" => {
if let Some(name) = graph::active_agent_graph_name(ctx) {
bail!(
"Graph-based agent '{name}' does not support sessions. \
The graph manages its own state."
);
}
let app = Arc::clone(&ctx.app.config);
ctx.use_session(app.as_ref(), args, abort_signal.clone())
.await?;
@@ -498,13 +531,41 @@ pub async fn run_repl_command(
};
eprintln!("\n📢 {}", color.italic().paint("Autonaming the session."),);
if let Err(err) = ctx.autoname_session(app.as_ref()).await {
log::warn!("Failed to autonaming the session: {err}");
warn!("Failed to autonaming the session: {err}");
}
if let Some(session) = ctx.session.as_mut() {
session.set_autonaming(false);
}
}
}
".install" => {
let trimmed = args.map(str::trim).unwrap_or("");
let mut parts = trimmed.splitn(2, char::is_whitespace);
match parts.next() {
Some("remote") => {
let rest = parts.next().unwrap_or("").trim();
config::install_remote_from_repl_args(rest)?;
}
Some(name) if !name.is_empty() => match AssetCategory::parse(name) {
Some(category) => config::install_assets(category)?,
None => println!(
"Unknown asset category '{name}'. Valid categories: {}",
AssetCategory::NAMES.join(", ")
),
},
_ => println!(
"Usage: .install <{}> | .install remote <git-url>",
AssetCategory::NAMES.join("|")
),
}
}
".update" => {
if ctx.macro_flag {
bail!("Cannot perform this operation because you are in a macro")
}
let version = args.map(|s| s.trim().to_string());
task::spawn_blocking(move || config::run_self_update(version, false)).await??;
}
".rag" => {
ctx.use_rag(args, abort_signal.clone()).await?;
}
@@ -595,8 +656,13 @@ pub async fn run_repl_command(
let app = Arc::clone(&ctx.app.config);
ctx.edit_agent_config(app.as_ref())?;
}
Some("mcp-config") => {
ctx.edit_mcp_config()?;
}
_ => {
println!(r#"Usage: .edit <config|role|session|rag-docs|agent-config>"#)
println!(
r#"Usage: .edit <config|mcp-config|role|session|rag-docs|agent-config>"#
)
}
}
}
@@ -764,26 +830,19 @@ pub async fn run_repl_command(
bail!("Use '.empty session' instead");
}
Some("todo") => {
let cleared = match ctx.agent.as_mut() {
Some(agent) => {
if !agent.auto_continue_enabled() {
let config = ctx.auto_continue_config();
if !config.enabled {
bail!(
"The todo system is not enabled for this agent. Set 'auto_continue: true' in the agent's config.yaml to enable it."
"Auto-continue is not enabled. Set 'auto_continue: true' in your config to enable it."
);
}
if ctx.todo_list.is_empty() {
println!("Todo list is already empty.");
false
} else {
ctx.clear_todo_list();
println!("Todo list cleared.");
true
}
}
None => bail!("No active agent"),
};
let _ = cleared;
}
_ => unknown_command()?,
},
".vault" => match split_first_arg(args) {
@@ -855,8 +914,18 @@ async fn ask(
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
let client = input.create_client()?;
let app = Arc::clone(&ctx.app.config);
if graph::active_agent_graph_name(ctx).is_some() {
ctx.before_chat_completion(&input)?;
let output =
graph::run_active_agent_graph(ctx, &input.text(), abort_signal.clone()).await?;
app.print_markdown(&output)?;
ctx.after_chat_completion(app.as_ref(), &input, &output, &[])?;
return Ok(());
}
let client = input.create_client()?;
ctx.before_chat_completion(&input)?;
let (output, tool_results) = if input.stream() {
call_chat_completions_streaming(&input, client.as_ref(), ctx, abort_signal.clone()).await?
@@ -881,19 +950,22 @@ async fn ask(
)
.await
} else {
let should_continue = agent_should_continue(ctx);
let do_continue = should_continue(ctx);
if should_continue {
if do_continue {
let full_prompt = {
let config = ctx.auto_continue_config();
let todo_state = ctx.todo_list.render_for_model();
let remaining = ctx.todo_list.incomplete_count();
ctx.set_last_continuation_response(output.clone());
ctx.increment_auto_continue_count();
let agent = ctx.agent.as_mut().expect("agent checked above");
let count = ctx.auto_continue_count;
let max = agent.max_auto_continues();
let max = config.max_continues;
let prompt = agent.continuation_prompt();
let prompt = config
.continuation_prompt
.as_deref()
.unwrap_or(DEFAULT_CONTINUATION_PROMPT);
let color = if app.light_theme() {
nu_ansi_term::Color::LightGray
@@ -921,7 +993,7 @@ async fn ask(
};
eprintln!("\n📢 {}", color.italic().paint("Autonaming the session."),);
if let Err(err) = ctx.autoname_session(app.as_ref()).await {
log::warn!("Failed to autonaming the session: {err}");
warn!("Failed to autonaming the session: {err}");
}
if let Some(session) = ctx.session.as_mut() {
session.set_autonaming(false);
@@ -934,7 +1006,7 @@ async fn ask(
.is_some_and(|s| s.needs_compression(app.compression_threshold));
if needs_compression {
let agent_can_continue_after_compress = agent_should_continue(ctx);
let agent_can_continue_after_compress = should_continue(ctx);
if let Some(session) = ctx.session.as_mut() {
session.set_compressing(true);
@@ -948,7 +1020,7 @@ async fn ask(
eprintln!("\n📢 {}", color.italic().paint("Compressing the session."),);
if let Err(err) = ctx.compress_session().await {
log::warn!("Failed to compress the session: {err}");
warn!("Failed to compress the session: {err}");
}
if let Some(session) = ctx.session.as_mut() {
session.set_compressing(false);
@@ -956,14 +1028,17 @@ async fn ask(
if agent_can_continue_after_compress {
let full_prompt = {
let config = ctx.auto_continue_config();
let todo_state = ctx.todo_list.render_for_model();
let remaining = ctx.todo_list.incomplete_count();
ctx.increment_auto_continue_count();
let agent = ctx.agent.as_mut().expect("agent checked above");
let count = ctx.auto_continue_count;
let max = agent.max_auto_continues();
let max = config.max_continues;
let prompt = agent.continuation_prompt();
let prompt = config
.continuation_prompt
.as_deref()
.unwrap_or(DEFAULT_CONTINUATION_PROMPT);
let color = if app.light_theme() {
nu_ansi_term::Color::LightGray
@@ -989,10 +1064,12 @@ async fn ask(
}
}
fn agent_should_continue(ctx: &RequestContext) -> bool {
ctx.agent.as_ref().is_some_and(|agent| {
agent.auto_continue_enabled() && ctx.auto_continue_count < agent.max_auto_continues()
}) && ctx.todo_list.has_incomplete()
fn should_continue(ctx: &RequestContext) -> bool {
let config = ctx.auto_continue_config();
ctx.app.config.function_calling_support
&& config.enabled
&& ctx.auto_continue_count < config.max_continues
&& ctx.todo_list.has_incomplete()
}
fn reset_continuation(ctx: &mut RequestContext) {
@@ -1188,8 +1265,8 @@ mod tests {
}
#[test]
fn repl_commands_has_39_entries() {
assert_eq!(REPL_COMMANDS.len(), 39);
fn repl_commands_has_42_entries() {
assert_eq!(REPL_COMMANDS.len(), 42);
}
#[test]
@@ -1311,13 +1388,15 @@ mod tests {
}
#[test]
fn repl_commands_clear_todo_requires_agent() {
fn repl_commands_clear_todo_always_available() {
let cmd = REPL_COMMANDS
.iter()
.find(|c| c.name == ".clear todo")
.unwrap();
assert!(cmd.is_valid(StateFlags::AGENT));
assert!(!cmd.is_valid(StateFlags::empty()));
assert!(cmd.is_valid(StateFlags::empty()));
assert!(cmd.is_valid(StateFlags::SESSION));
assert!(cmd.is_valid(StateFlags::ROLE));
}
#[test]
+2 -16
View File
@@ -34,7 +34,6 @@ use is_terminal::IsTerminal;
use std::borrow::Cow;
use std::sync::LazyLock;
use std::{cmp, env, path::PathBuf, process};
use unicode_segmentation::UnicodeSegmentation;
pub static CODE_BLOCK_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?ms)```\w*(.*)```").unwrap());
@@ -74,21 +73,8 @@ pub fn parse_bool(value: &str) -> Option<bool> {
}
pub fn estimate_token_length(text: &str) -> usize {
let words: Vec<&str> = text.unicode_words().collect();
let mut output: f32 = 0.0;
for word in words {
if word.is_ascii() {
output += 1.3;
} else {
let count = word.chars().count();
if count == 1 {
output += 1.0
} else {
output += (count as f32) * 0.5;
}
}
}
output.ceil() as usize
let weighted: usize = text.chars().map(|c| if c.is_ascii() { 1 } else { 2 }).sum();
weighted.div_ceil(4)
}
pub fn strip_think_tag(text: &str) -> Cow<'_, str> {
+54 -15
View File
@@ -241,23 +241,23 @@ fn add_file(files: &mut IndexSet<String>, suffixes: Option<&Vec<String>>, path:
}
fn is_valid_extension(suffixes: Option<&Vec<String>>, path: &Path) -> bool {
let filename_regex = Regex::new(r"^.+\.*").unwrap();
if let Some(suffixes) = suffixes
&& !suffixes.is_empty()
{
if let Ok(Some(_)) = filename_regex.find(&suffixes.join(",")) {
let file_name = path
.file_name()
.and_then(|v| v.to_str())
.expect("invalid filename")
.to_string();
return suffixes.contains(&file_name);
} else if let Some(extension) = path.extension().map(|v| v.to_string_lossy().to_string()) {
return suffixes.contains(&extension);
let Some(suffixes) = suffixes else {
return true;
};
if suffixes.is_empty() {
return true;
}
return false;
let file_name = path.file_name().and_then(|v| v.to_str());
let extension = path.extension().and_then(|v| v.to_str());
suffixes.iter().any(|suffix| {
if suffix.contains('.') {
Some(suffix.as_str()) == file_name
} else {
Some(suffix.as_str()) == extension
}
true
})
}
#[cfg(test)]
@@ -352,4 +352,43 @@ mod tests {
)
);
}
#[test]
fn test_is_valid_extension() {
let md_ext = vec!["md".to_string()];
let md_txt_ext = vec!["md".to_string(), "txt".to_string()];
let test_md_filename = vec!["test.md".to_string()];
let mixed = vec!["md".to_string(), "test.txt".to_string()];
assert!(is_valid_extension(None, Path::new("Agents.md")));
assert!(is_valid_extension(Some(&vec![]), Path::new("Agents.md")));
assert!(is_valid_extension(Some(&md_ext), Path::new("Agents.md")));
assert!(is_valid_extension(
Some(&md_ext),
Path::new("/home/atusa/code/loki.wiki/Agents.md")
));
assert!(!is_valid_extension(Some(&md_ext), Path::new("notes.txt")));
assert!(!is_valid_extension(Some(&md_ext), Path::new("README")));
assert!(is_valid_extension(Some(&md_txt_ext), Path::new("a.md")));
assert!(is_valid_extension(Some(&md_txt_ext), Path::new("a.txt")));
assert!(!is_valid_extension(Some(&md_txt_ext), Path::new("a.rs")));
assert!(is_valid_extension(
Some(&test_md_filename),
Path::new("dir/test.md")
));
assert!(!is_valid_extension(
Some(&test_md_filename),
Path::new("dir/Agents.md")
));
assert!(is_valid_extension(Some(&mixed), Path::new("Agents.md")));
assert!(is_valid_extension(Some(&mixed), Path::new("dir/test.txt")));
assert!(!is_valid_extension(
Some(&mixed),
Path::new("dir/other.txt")
));
}
}