refactor(runtime): reconcile unified kernel flow

Align the active context, memory, and tool runtime paths with the shipped kernel so the branch reflects the real production execution model.
Capture the final verification baseline in code and docs, including the last eval hardening fixes that brought the full suite back to green.
This commit is contained in:
ZanzyTHEbar 2026-03-22 16:33:52 +00:00
parent 9313f141fe
commit 1bdd82cc47
80 changed files with 6607 additions and 431 deletions

View file

@ -118,15 +118,15 @@ flowchart TB
| Decision | Rationale | | Decision | Rationale |
|----------|-----------| |----------|-----------|
| **Isolated Tool Runtime** | All tool calls route through a `SecureBus` that enforces capability manifests, injects secrets, scans output for leaks, and writes audit logs. The LLM never sees raw secrets. See [ADR-001](docs/adr/001-isolated-tool-runtime.md). | | **Isolated Tool Runtime** | All tool calls route through a `SecureBus` that enforces capability manifests, injects secrets, scans output for leaks, and writes audit logs. The LLM never sees raw secrets. Layer 1-2 are live today; daemon/WASM isolation remains optional roadmap work. See [ADR-001](docs/adr/001-isolated-tool-runtime.md). |
| **DAG executor** | LLMCompiler-style parallel tool dispatch. The planner builds a dependency DAG in a single inference pass; the executor dispatches independent nodes concurrently. Joiner synthesizes results. Replanning on failure. Falls back to ReAct for simple single-tool cases. | | **DAG executor** | LLMCompiler-style parallel tool dispatch. The planner builds a dependency DAG in a single inference pass; the executor dispatches independent nodes concurrently. Joiner synthesizes results. Replanning on failure. Falls back to ReAct for simple single-tool cases. |
| **Vendored Fantasy SDK** | `charm.land/fantasy` vendored into `internal/fantasy/` via `go.mod` replace directive. Enables direct modification for streaming hooks, tool call repair, and progressive disclosure. | | **Vendored Fantasy SDK** | `charm.land/fantasy` vendored into `internal/fantasy/` via `go.mod` replace directive. Enables direct modification for streaming hooks, tool call repair, and progressive disclosure. |
| **MemGPT + Observational Memory** | Working context (hot), recall items (warm), archival chunks (cold, embedded + indexed), plus observational memory (compressed conversation history with priority-tagged observations and temporal reasoning). DAG-based context budget compression manages token allocation across tiers. | | **MemGPT + Projection Kernel** | Working context (hot), recall items (warm), archival chunks (cold, embedded + indexed), observational memory, immutable messages, DAG snapshots, runtime checkpoints, and an active-context projection builder that assembles the live turn context. Semantic ContextTree scoring and RLM reduction now participate in the hot path. |
| **Progressive tool disclosure** | Agent sees only `tool_search` and `tool_call` meta-tools. Discovers actual tools on demand via fuzzy search. Cuts system prompt tokens for large registries. | | **Progressive tool disclosure** | Agent sees only `tool_search` and `tool_call` meta-tools. Discovers actual tools on demand via fuzzy search. Cuts system prompt tokens for large registries. |
| **libSQL over modernc/sqlite** | Native F32_BLOB for vector storage, `libsql_vector_idx` for ANN search, FTS5 for full-text. Single database, no external vector DB dependency. | | **libSQL over modernc/sqlite** | Native F32_BLOB for vector storage, `libsql_vector_idx` for ANN search, FTS5 for full-text. Single database, no external vector DB dependency. |
| **BLOB primary keys** | 16-byte UUIDv7 stored as BLOB. Compact, byte-comparable, monotonically sortable by creation time. | | **BLOB primary keys** | 16-byte UUIDv7 stored as BLOB. Compact, byte-comparable, monotonically sortable by creation time. |
| **XChaCha20-Poly1305 vault** | Secrets encrypted at rest with AES-256-GCM or XChaCha20-Poly1305. Master key from OS keyring, env var, or file. Schnorr ZKP for daemon-mode authentication. | | **XChaCha20-Poly1305 vault** | Secrets encrypted at rest with AES-256-GCM or XChaCha20-Poly1305. Master key from OS keyring, env var, or file. Schnorr ZKP remains planned for daemon-mode authentication. |
| **Goose migrations** | Schema managed by `pressly/goose/v3`. 10 versioned migrations covering core schema, FTS5, vector indexes, KV store, documents, audit log, conversations, runtime state, jobs, and conversation graphs. | | **Goose migrations** | Schema managed by `pressly/goose/v3`. 17 versioned migrations currently cover core schema, FTS5, vector indexes, KV store, documents, audit log outcomes, conversations, runtime state, DAG/checkpoint data, map operators, memory edges, soft delete, immutable messages, and RL/task-completion tables. |
| **FlatBuffers command protocol** | Zero-copy serialized `ToolRequest`/`ToolResponse` for the ITR command vocabulary. Same binary format across in-process channels, Unix sockets (daemon mode), and wazero WASM host calls. | | **FlatBuffers command protocol** | Zero-copy serialized `ToolRequest`/`ToolResponse` for the ITR command vocabulary. Same binary format across in-process channels, Unix sockets (daemon mode), and wazero WASM host calls. |
## Project Layout ## Project Layout
@ -220,7 +220,6 @@ Edit `~/.dragonscale/config.json`:
} }
}, },
"tools": { "tools": {
"progressive_disclosure": true,
"web": { "web": {
"duckduckgo": { "enabled": true, "max_results": 5 } "duckduckgo": { "enabled": true, "max_results": 5 }
} }
@ -422,7 +421,7 @@ API key links: [OpenRouter](https://openrouter.ai/keys) · [Anthropic](https://c
## Memory System ## Memory System
DragonScale implements a multi-tier memory system combining MemGPT-style tiered storage with observational memory compression: DragonScale implements a multi-tier memory system combining MemGPT-style storage with immutable history, DAG snapshots, semantic context selection, and active-context projection:
| Tier | Purpose | Storage | Search | | Tier | Purpose | Storage | Search |
|------|---------|---------|--------| |------|---------|---------|--------|
@ -430,11 +429,11 @@ DragonScale implements a multi-tier memory system combining MemGPT-style tiered
| **Recall Memory** | Session-scoped conversation items | Rows with metadata + timestamps | FTS5 + BM25 | | **Recall Memory** | Session-scoped conversation items | Rows with metadata + timestamps | FTS5 + BM25 |
| **Archival Memory** | Long-term knowledge, chunked + embedded | F32_BLOB embeddings + FTS5 index | Vector ANN + FTS5 fusion (RRF) | | **Archival Memory** | Long-term knowledge, chunked + embedded | F32_BLOB embeddings + FTS5 index | Vector ANN + FTS5 fusion (RRF) |
| **Observational Memory** | Compressed conversation history | Priority-tagged observations with 3-date model | Prefix-cacheable block | | **Observational Memory** | Compressed conversation history | Priority-tagged observations with 3-date model | Prefix-cacheable block |
| **DAG Compression** | Hierarchical context summaries | Tree nodes with lossless pointers to originals | Budget-allocated traversal | | **DAG Snapshots** | Persistent hierarchical summaries + recovery | Snapshot/node/edge tables with lossless refs | Projection tier + DAG tools |
The agent interacts with memory through a unified `memory` tool. Large tool results are automatically offloaded to archival memory. Retrieval uses Reciprocal Rank Fusion (RRF) to combine vector similarity and full-text relevance, with recency decay and metadata pre-filtering. The agent interacts with memory through a unified `memory` tool. Large tool results are automatically offloaded to archival memory. Retrieval uses Reciprocal Rank Fusion (RRF) to combine vector similarity and full-text relevance, with recency decay and metadata pre-filtering. Turn assembly now flows through `ActiveContextProjection`, which can combine recent immutable history, semantic ContextTree-selected recall, DAG projection segments, and RLM-reduced long-context slices before prompt rendering.
Schema is managed by Goose with 10 versioned migrations. Schema is managed by Goose with 17 versioned migrations.
## CLI Reference ## CLI Reference
@ -468,7 +467,7 @@ The SecureBus mediates all tool execution. Tools declare capabilities via the `C
The pipeline: capability check → secret injection → tool execution → leak scanning → audit log. The pipeline: capability check → secret injection → tool execution → leak scanning → audit log.
See [ADR-001](docs/adr/001-isolated-tool-runtime.md) for the full design including DAG executor convergence, RLM integration, and the FlatBuffers command protocol. See [ADR-001](docs/adr/001-isolated-tool-runtime.md) for the layered design, current rollout status, and the FlatBuffers command protocol.
## Development ## Development
@ -550,6 +549,9 @@ A promptfoo-based evaluation harness lives in `eval/`:
cd eval && go run ./cmd/eval-runner cd eval && go run ./cmd/eval-runner
``` ```
`make eval` auto-builds `eval/bin/eval-runner` when it is missing.
`make eval-compare` builds the comparison binary for `main` in a temporary git worktree so the current checkout is never stashed or switched underneath the user.
### Syncing Upstream ### Syncing Upstream
```bash ```bash

View file

@ -12,8 +12,14 @@ Reference blueprint: `docs/execution/unified-kernel-blueprint.md`
- [x] Single always-on runtime path (SecureBus + offloading + run-state persistence). - [x] Single always-on runtime path (SecureBus + offloading + run-state persistence).
- [x] Fail-fast boot invariants for kernel dependencies. - [x] Fail-fast boot invariants for kernel dependencies.
- [x] Live ReAct transition persistence, authoritative step indexing, and true task tool-call metrics.
- [x] Deterministic session continuity with projection pointers + integrity validation. - [x] Deterministic session continuity with projection pointers + integrity validation.
- [x] Emergency-only recursive compression and provenance persistence. - [x] Emergency-only recursive compression and provenance persistence.
- [x] Active-context projection builder now owns hot-path assembly across system, recent history, retrieval, and DAG tiers.
- [x] Session-correct working context + memory tool binding + durable session summaries.
- [x] Semantic ContextTree scoring, preserved access state, and DAG projection segments are active in production.
- [x] Runtime checkpoints persist and restore/fork sessions through the existing checkpoint schema.
- [x] Baseline RLM reduction is wired into production context assembly for oversized DAG/recall/archival segments.
- [x] Persistent DAG snapshots + lossless DAG tools (`dag_expand`, `dag_describe`, `dag_grep`). - [x] Persistent DAG snapshots + lossless DAG tools (`dag_expand`, `dag_describe`, `dag_grep`).
- [x] Subagent runtime parity with delegation scope/lineage/depth/fanout guardrails. - [x] Subagent runtime parity with delegation scope/lineage/depth/fanout guardrails.
- [x] Assistant-focused proactive eval suite (commitments/reminders/follow-ups/continuity). - [x] Assistant-focused proactive eval suite (commitments/reminders/follow-ups/continuity).
@ -22,6 +28,7 @@ Reference blueprint: `docs/execution/unified-kernel-blueprint.md`
- [x] Map operators are now FlatBuffers-first end to end for persisted run/item state (`spec_fb`, `input_fb`, `output_fb`). - [x] Map operators are now FlatBuffers-first end to end for persisted run/item state (`spec_fb`, `input_fb`, `output_fb`).
- [x] Worker identity and deduplication for map jobs now resolve via deterministic keys (`map:{runID}:{itemIndex}`), with idempotent run reuse under concurrency. - [x] Worker identity and deduplication for map jobs now resolve via deterministic keys (`map:{runID}:{itemIndex}`), with idempotent run reuse under concurrency.
- [x] Devcontainer + `flatc` + `sqlc` generation pipeline is active (`make devcontainer-build`, `make devcontainer-up`, `make devcontainer-generate`, `make devcontainer-verify`). - [x] Devcontainer + `flatc` + `sqlc` generation pipeline is active (`make devcontainer-build`, `make devcontainer-up`, `make devcontainer-generate`, `make devcontainer-verify`).
- [x] Audit rows now persist explicit outcome fields (`success`, `error_msg`, `tool_call_id`) and the eval harness enforces strict default assertions.
--- ---
@ -120,13 +127,15 @@ Deterministic engine compresses old messages into a hierarchical DAG while keepi
### 5d. RLM Memory Controller ### 5d. RLM Memory Controller
*Dedicated controller layer for Memory I/O — optimizes token usage, call frequency, read/write scheduling.* *Baseline recursive reduction is live today.
The remaining roadmap work is a fuller controller layer for coalesced reads, buffered writes, and adaptive memory I/O policy.*
- **Production baseline (shipped)** — Active-context assembly can route oversized DAG / recall / archival projection segments through `pkg/rlm` before prompt rendering.
- **Read coalescing** — Batch multiple memory reads into a single DB round-trip. Reduce I/O calls per agent loop iteration. - **Read coalescing** — Batch multiple memory reads into a single DB round-trip. Reduce I/O calls per agent loop iteration.
- **Write buffering** — Buffer memory writes and flush on consolidation boundaries (aligned with Focus checkpoints). - **Write buffering** — Buffer memory writes and flush on consolidation boundaries (aligned with Focus checkpoints).
- **Token budget enforcement** — Hard cap on tokens loaded from memory per turn. Controller decides what to fetch given the budget, using importance scores from the existing scoring pipeline. - **Token budget enforcement** — Hard cap on tokens loaded from memory per turn. Controller decides what to fetch given the budget, using importance scores from the existing scoring pipeline.
- **Adaptive fetch** — Controller adjusts retrieval depth based on task complexity signal (number of tool calls, error rate, context pressure). - **Adaptive fetch** — Controller adjusts retrieval depth based on task complexity signal (number of tool calls, error rate, context pressure).
- **Recursive Language Models**: Utilize a full in-process RLM agent system to manage the memory store. - **Recursive DAG expansion** — Promote the current reducer into a deeper controller that can recurse over partitions and drive memory-heavy sub-queries end-to-end.
*Extends*: `pkg/memory/store/` — new `Controller` layer wrapping `MemoryStore`. *Extends*: `pkg/memory/store/` — new `Controller` layer wrapping `MemoryStore`.
@ -307,12 +316,13 @@ flowchart LR
- [ ] Tools can then live in packages like `pkg/tools/filesystem.go` and `pkg/tools/network.go` where we expose only a limited ABI/API to the agent loop - [ ] Tools can then live in packages like `pkg/tools/filesystem.go` and `pkg/tools/network.go` where we expose only a limited ABI/API to the agent loop
- [ ] Extract out all inline prompts into separate files - [ ] Extract out all inline prompts into separate files
- [ ] Use dotprompt, poml, or any other structured prompt builder to build prompts - [ ] Use dotprompt, poml, or any other structured prompt builder to build prompts
- [ ] Isolated tool runtime + DAG executor + RLM engine — see [ADR-001](docs/adr/001-isolated-tool-runtime.md) - [ ] Further expansion of isolated tool runtime + DAG executor + RLM engine — see [ADR-001](docs/adr/001-isolated-tool-runtime.md)
- [ ] Layer 1: Capability manifests (`CapableTool` interface) - [x] Layer 1: Capability manifests (`CapableTool` interface)
- [ ] Layer 2: SecureBus + FlatBuffers command protocol (incl. DAG types) + leak scanning - [x] Layer 2: SecureBus + FlatBuffers command protocol (incl. DAG types) + leak scanning
- [ ] DAG executor: LLMCompiler-style parallel dispatch, topological wave execution, dependency resolution, Joiner synthesis, replanning loop - [x] DAG executor: dependency-aware parallel dispatch is active on the tool runtime path
- [ ] Programmatic tool calling: PTC-style context isolation (intermediate results never enter LLM context), ToolSearch for on-demand tool discovery - [x] Programmatic tool calling baseline: intermediate tool results are offloaded/indexed and only compact previews enter prompt assembly
- [ ] RLM engine: recursive context decomposition (rope DS, parallel fan-out, cheap sub-LM strategy, recursive DAG expansion) - [x] RLM baseline: production context reduction over oversized projection segments
- [ ] Full recursive DAG expansion and deeper memory-controller orchestration
- [ ] ReAct/DAG routing: automatic mode selection (`ModeReAct | ModeDAG | ModeAuto`) - [ ] ReAct/DAG routing: automatic mode selection (`ModeReAct | ModeDAG | ModeAuto`)
- [ ] Layer 3: SecretStore + keyring-based secret management - [ ] Layer 3: SecretStore + keyring-based secret management
- [ ] Layer 4: Daemon mode + Schnorr ZKP authentication - [ ] Layer 4: Daemon mode + Schnorr ZKP authentication

View file

@ -1,11 +1,29 @@
# ADR-001: Isolated Tool Runtime (ITR) + DAG Task Executor # ADR-001: Isolated Tool Runtime (ITR) + DAG Task Executor
**Date**: 2026-02-18 (updated 2026-02-19) **Date**: 2026-02-18 (updated 2026-02-19)
**Status**: Proposed **Status**: Accepted (incremental rollout)
**Authors**: @ZanzyTHEbar **Authors**: @ZanzyTHEbar
--- ---
## Current Rollout Status
This ADR mixes shipped kernel behavior with the longer-range secure execution roadmap.
### Live in production
- Layer 1-2 SecureBus mediation is active for tool execution.
- FlatBuffers command vocabulary is live for the internal command surface.
- Capability enforcement, secret injection, leak scanning, and audit logging are on the hot path.
- Dependency-aware parallel tool execution is active through the DAG tool runtime.
- `pkg/rlm` is now wired into active-context assembly as a reducer for oversized DAG / recall / archival projection segments.
### Still planned
- Daemon-mode SecureBus separation and Schnorr ZKP session establishment.
- wazero-backed isolated execution for untrusted tools.
- Full DAG-planner-time recursive RLM expansion through SecureBus as the universal long-context backbone.
## Context ## Context
DragonScale tools execute in-process with the agent loop. The `Vault` (XChaCha20-Poly1305) exists for encrypting secrets at rest, but there is no pipeline for injecting those secrets into tool execution. The `Redactor` scans for sensitive patterns, but only in log paths — not on tool output before it reaches the LLM. There is no privilege boundary DragonScale tools execute in-process with the agent loop. The `Vault` (XChaCha20-Poly1305) exists for encrypting secrets at rest, but there is no pipeline for injecting those secrets into tool execution. The `Redactor` scans for sensitive patterns, but only in log paths — not on tool output before it reaches the LLM. There is no privilege boundary

View file

@ -56,5 +56,10 @@ Additional decisions:
- Core wiring is in `pkg/agent/loop.go`, `pkg/agent/securebus_runtime.go`, and `pkg/runtime/bootstrap.go`. - Core wiring is in `pkg/agent/loop.go`, `pkg/agent/securebus_runtime.go`, and `pkg/runtime/bootstrap.go`.
- Session continuity changes are in `pkg/session/manager.go`. - Session continuity changes are in `pkg/session/manager.go`.
- Live transition persistence, run-state bookkeeping, and runtime checkpoints are wired through `pkg/agent/agent_run.go`, `pkg/agent/checkpoint_runtime.go`, and `pkg/agent/state_store.go`.
- Active-context assembly is now owned by `pkg/agent/active_context_builder.go` and rendered by `pkg/agent/context.go`.
- Semantic ContextTree scoring and DAG projection segments are active in `pkg/agent/summarizer.go` and `pkg/agent/active_context_builder.go`.
- Baseline RLM reduction is active in `pkg/agent/rlm_runtime.go` as part of the context-assembly hot path.
- Explicit audit outcome persistence (`success`, `error_msg`, `tool_call_id`) is active in `pkg/memory/delegate/sqlite.go` and `pkg/memory/migrations/017_agent_audit_outcomes.go`.
- Subagent control-flow parity improvements are in `pkg/agent/toolloop.go` and `pkg/tools/subagent.go`. - Subagent control-flow parity improvements are in `pkg/agent/toolloop.go` and `pkg/tools/subagent.go`.

View file

@ -0,0 +1,174 @@
# Hermes Alignment for DragonScale
## Status
Frozen for future work.
## Purpose
Capture the lasting strategic takeaways from the Hermes research pass so DragonScale can borrow the right ideas without weakening its kernel architecture.
This document is not an implementation spec.
Linear remains the source of truth for execution.
Memory Bank remains the source of truth for durable strategic context.
## Strategic Thesis
DragonScale should not become "Hermes in Go."
DragonScale already has the stronger context kernel:
- projection-first turn assembly
- immutable history with lossless references
- DAG-backed recovery and retrieval
- deterministic unified runtime invariants
- RLM reduction integrated into the hot path
Hermes is stronger in a different layer:
- runtime ergonomics
- prompt economics
- bounded always-on memory for stable preferences
- skill self-maintenance
- operator-facing product shell
The right synthesis is:
- keep DragonScale's kernel contracts
- borrow Hermes-grade ergonomics around the kernel
## What To Keep
- `ActiveContextProjection` as the canonical turn contract
- immutable history and lossless recovery references
- unified runtime path from `docs/adr/002-unified-kernel-runtime.md`
- emergency-safe context reduction and budgeted assembly
- clear separation between deterministic kernel behavior and agent autonomy
## What To Borrow
### 1. Prompt stability as an invariant
Adopt more explicit stable vs session-stable vs turn-volatile prompt segmentation.
Goal:
- increase provider-side prompt cache reuse
- reduce unnecessary prompt rebuilding
- preserve the projection kernel while making rendering cheaper and more legible
### 2. Tiny bounded profile memory
Add a small always-on profile layer for durable user preferences and agent operating constraints.
This should be:
- size capped
- prompt-native rather than retrieval-first
- separate from recall and archival memory
### 3. First-class session recall surface
Expose session-scoped history retrieval as a first-class product surface.
Goal:
- easy "what did we already do?" recall
- lineage-aware history lookup
- cleaner operator and agent access to prior work
### 4. Skill maintenance loop
Treat skills as living procedural memory rather than static markdown assets.
Goal:
- detect stale skills
- propose or apply fixes under policy
- track skill freshness and evidence of successful reuse
### 5. Explicit auxiliary model roles
Make model-role boundaries more explicit in config and runtime documentation.
Examples:
- primary task model
- compression model
- cheap utility model
- extraction model
- fallback model
### 6. Pressure signaling
Expose context and iteration pressure more clearly to both the agent and the operator.
Goal:
- encourage consolidation before failure
- make runtime state legible instead of implicit
### 7. Operator shell and inspection surfaces
Improve observability with explicit product-shell commands and inspections.
Examples:
- `doctor`
- `config check`
- `context inspect`
- `projection inspect`
- `session inspect`
- `retrieval policy status`
## What Not To Copy
- file-backed markdown memory as the primary durable memory system
- transcript-first compression as the primary context engine
- broad agent autonomy over core memory mechanics
- monolithic loop design that weakens kernel boundaries
## Current Overlap With Existing Work
Some Hermes-adjacent ideas already exist in the roadmap and backlog:
- stable-prefix context ideas overlap with Observational Memory
- retrieval surface work overlaps with Agentic Retrieval
- model-role ideas overlap partially with SubAgent Profiles and model routing
- CLI/inspection improvements overlap partially with the Cobra migration
The Hermes pass should therefore create:
- one parent strategy item
- a small number of missing child issues
- explicit references to overlapping work
It should not duplicate existing roadmap items.
## New Work That Appears Missing
The research pass identified these gaps as not yet cleanly represented:
1. bounded micro-profile memory
2. explicit session-search / history-search product surface
3. skill maintenance and self-patching loop
4. explicit auxiliary model-role architecture
5. context and iteration pressure signaling
6. operator diagnostics and inspection shell
## Decision Rule
Future Hermes-inspired work must satisfy both constraints:
1. It improves operator ergonomics, prompt economics, or bounded memory behavior.
2. It does not weaken the projection-first kernel, immutable history model, or unified runtime invariants.
If a proposed Hermes-inspired change conflicts with kernel determinism, DragonScale keeps the kernel and rejects the transplant.
## References
- `README.md`
- `ROADMAP.md`
- `docs/adr/002-unified-kernel-runtime.md`
- `docs/execution/unified-kernel-blueprint.md`
- [Hermes research](956e0452-1f29-4344-ad93-2cea783ae5a6)

View file

@ -17,18 +17,36 @@ This blueprint defines the always-on runtime model for DragonScale's assistant-f
- Unified runtime assembly is active in `pkg/agent/loop.go` and `pkg/agent/securebus_runtime.go`. - Unified runtime assembly is active in `pkg/agent/loop.go` and `pkg/agent/securebus_runtime.go`.
- Bootstrap fail-fast checks are enforced in `pkg/runtime/bootstrap.go`. - Bootstrap fail-fast checks are enforced in `pkg/runtime/bootstrap.go`.
- ReAct transition persistence, authoritative step indexing, and corrected task metrics are active in `pkg/agent/agent_run.go`.
- Session projection pointers + integrity validation are active in `pkg/session/manager.go` and `pkg/session/projection_pointer.go`. - Session projection pointers + integrity validation are active in `pkg/session/manager.go` and `pkg/session/projection_pointer.go`.
- Session-correct working context binding and durable session summaries are active in `pkg/agent/context.go`, `pkg/agent/memgpt_tool.go`, and `pkg/session/manager.go`.
- Active-context projection assembly is live in `pkg/agent/active_context_builder.go`.
- Emergency compression provenance capture is active in `pkg/agent/loop.go`. - Emergency compression provenance capture is active in `pkg/agent/loop.go`.
- DAG persistence and retrieval tools are active in: - DAG persistence and retrieval tools are active in:
- `pkg/memory/dag/store.go` - `pkg/memory/dag/store.go`
- `pkg/tools/dag.go` - `pkg/tools/dag.go`
- `pkg/memory/sqlc/queries/dag.sql` - `pkg/memory/sqlc/queries/dag.sql`
- Semantic ContextTree scoring, preserved access state, and DAG projection segments are active in:
- `pkg/agent/summarizer.go`
- `pkg/contexttree/tree.go`
- `pkg/agent/active_context_builder.go`
- Runtime checkpoints and session restore/fork hydration are active in:
- `pkg/agent/checkpoint_runtime.go`
- `pkg/agent/conversations/checkpoint_snapshot.go`
- `pkg/agent/conversations/store.go`
- Baseline RLM reduction is active in:
- `pkg/agent/rlm_runtime.go`
- `pkg/agent/agent_run.go`
- Legacy session and DAG backfill passes are active at startup: - Legacy session and DAG backfill passes are active at startup:
- Session pointer backfill status: `migration:session_projection_backfill:v1` - Session pointer backfill status: `migration:session_projection_backfill:v1`
- DAG backfill status: `migration:dag_backfill:v1` - DAG backfill status: `migration:dag_backfill:v1`
- Subagent delegation safety (scope/kept-work, depth/fanout, lineage audit) is active in `pkg/tools/subagent.go`. - Subagent delegation safety (scope/kept-work, depth/fanout, lineage audit) is active in `pkg/tools/subagent.go`.
- Hybrid retrieval routing across working-context, recall, archival, and DAG projections is active in `pkg/memory/store/memory_store.go`. - Hybrid retrieval routing across working-context, recall, archival, and DAG projections is active in `pkg/memory/store/memory_store.go`.
- Shadow-mode rollout, proof gates, auto-promotion, and fast rollback for retrieval augmentation are active in `pkg/memory/store/retrieval_policy.go`. - Shadow-mode rollout, proof gates, auto-promotion, and fast rollback for retrieval augmentation are active in `pkg/memory/store/retrieval_policy.go`.
- Explicit audit outcome persistence is active in:
- `pkg/memory/migrations/017_agent_audit_outcomes.go`
- `pkg/memory/delegate/sqlite.go`
- `pkg/memory/sqlc/queries/agent_audit_log.sql`
- Map operator runtime is active with FlatBuffers persistence and worker orchestration in: - Map operator runtime is active with FlatBuffers persistence and worker orchestration in:
- `pkg/tools/map_runtime.go` - `pkg/tools/map_runtime.go`
- `pkg/tools/map_flatbuffer_codec.go` - `pkg/tools/map_flatbuffer_codec.go`
@ -49,4 +67,5 @@ This blueprint defines the always-on runtime model for DragonScale's assistant-f
## Remaining Work (Ordered) ## Remaining Work (Ordered)
- Integrate obligation heartbeat execution for proactive due checks. - Integrate obligation heartbeat execution for proactive due checks.
- Deepen the current RLM reducer into a fuller memory controller (read coalescing, write buffering, recursive partition orchestration).
- Keep JSONL strictly as an LLM boundary format; do not persist JSONL internally. - Keep JSONL strictly as an LLM boundary format; do not persist JSONL internally.

472
dragon-scale.md Normal file
View file

@ -0,0 +1,472 @@
# DragonScale Reconcile Report (Historical + Current)
## Audit Frame
- Scope: repository code, runtime architecture, memory stack, roadmap/docs/ADR drift, eval harness, and prior audit transcript context.
- Method: static reconciliation only. I inspected code and transcript artifacts; I did not run `go test`, `make eval`, TLC, or live external systems in this pass.
- Confidence: high on structural truth, medium on dynamic behavior that depends on live execution.
- Status note: sections `1` through `7` capture the pre-reconciliation audit state that drove the repair plan. The section `Reconciliation Status (2026-03-21)` below is the current implementation truth after the kernel-drift execution batches.
## Executive Truth
- The shipped core is a **unified kernel runtime**, not the older optional-path runtime. The real production stack is `Service.Agent|Gateway` -> `runtime.Bootstrap` -> `AgentLoop` -> `fantasy.Agent.Generate` -> `SecureBusToolRuntime` -> `OffloadingToolRuntime` -> `fantasy.DAGToolRuntime`.
- The memory system is materially real and fairly deep: session recall, immutable messages, archival chunks + embeddings, observational memory, DAG snapshots, focus knowledge blocks, projection pointers, hybrid retrieval, and background maintenance all exist in code.
- The repo currently mixes **three separate context systems**:
- `ContextBuilder` system-prompt assembly and section budgeting
- `ContextTree` query-adaptive selection over older history
- persistent `DAG` snapshots and DAG tools for lossless recovery/search
- The biggest remaining truth gaps are **observability/runtime bookkeeping**, not base execution:
- `agent_state_transitions` exists in schema but is not populated in production
- persisted/offloaded tool result `step_index` is not trustworthy
- `ActiveContextBuilder` exists as a contract only, not a live implementation
- `pkg/rlm` exists and is tested, but is not wired into production
- Documentation currently overstates some planned architecture and understates some shipped architecture.
## Reconciliation Status (2026-03-22)
- Runtime observability gaps called out above are now resolved in code: FSM transitions persist, step indices are propagated through tool execution, and task metrics count real tool activity instead of raw ReAct step count.
- The context kernel is now live as a production builder: `ActiveContextBuilder` assembles `ActiveContextProjection` segments, DAG projections are rendered into turn context, and session binding is resolved per active session rather than leaking through `default`.
- The memory kernel is materially stronger than the historical audit state: semantic ContextTree behavior is active, runtime checkpoints are created and can restore or fork sessions, and RLM reduction is wired into active-context assembly for oversized DAG / recall / archival segments.
- Audit persistence now stores explicit `success`, `error_msg`, and `tool_call_id` fields instead of inferring outcomes from action strings.
- Eval and docs reconciliation is complete for the kernel-drift scope: eval runner auto-build is live, `eval-compare` uses a temporary worktree, core docs (`README.md`, `ROADMAP.md`, ADRs, eval docs) were reconciled to implementation truth, and the final verification ladder is green.
- Remaining intentional gaps are now strictly future-work items:
- deeper RLM memory-controller orchestration beyond the current reducer baseline
- full daemon / WASM isolation stages from the longer-range secure-runtime roadmap
## Final Verification Evidence (2026-03-22)
- `go test ./...` passed.
- `go test -race ./...` passed.
- `make eval` passed at `62 passed / 0 failed / 0 errors` in `17m 8s`, writing the artifact to `eval/results/latest.json`.
- The last red eval was `memory search with no results: graceful empty response`; root cause was prompt-echo leakage from session-message recall rows and working-context projection during memory retrieval.
- Final remediation was applied in `pkg/memory/store/memory_store.go` by suppressing memory-search instruction echoes while keeping real stored matches, with regression coverage added in `pkg/memory/store/memory_tool_test.go`.
- Repository truth, report truth, and Memory Bank truth now converge on a fully green reconciliation closeout.
## Historical Findings (Pre-Reconciliation Baseline)
The next sections preserve the original audit findings that motivated the repair program.
Read them as historical counterexamples, not as the final current-state claim.
## Architectural Namespace Map
- Execution DAG: `internal/fantasy/tool_runtime_dag.go` for parallel tool execution from `$tool.<id>` dependencies.
- Memory DAG: `pkg/memory/dag/*` for deterministic hierarchical conversation compression and persistent snapshots.
- Skill Graph: `pkg/skills/*` + `pkg/tools/skills.go` for wikilink-based discovery and traversal.
- Context Tree: `pkg/contexttree/tree.go` for per-turn relevance scoring over historical messages.
- RLM: `pkg/rlm/*` exists as a separate decomposition engine but is not part of the current runtime path.
## 1. Control-Flow Truth
### Entry Modes
- `pkg/dragonscale/sdk/service_ops.go` has two real entry styles.
- CLI and eval-style direct calls use `Agent()` -> `ProcessDirect(...)`; this does **not** start `AgentLoop.Run`.
- Gateway mode uses `Gateway()` -> `go agentLoop.Run(appCtx)` and routes inbound channel traffic through `MessageBus`.
### Main Turn Path
- `pkg/agent/message_router.go` routes user traffic to `runAgentLoop(...)`.
- `pkg/agent/agent_run.go` assembles context, creates a per-turn `AgentConversation` and `AgentRun`, and builds the `fantasy.Agent`.
- `pkg/agent/agent_run.go` wires the runtime stack as:
- `SecureBusToolRuntime`
- wrapping `OffloadingToolRuntime`
- wrapping `fantasy.DAGToolRuntime`
- `internal/fantasy/agent.go` runs the ReAct loop:
- prepare step
- LLM call
- tool validation
- tool execution
- append messages
- stop check
- repeat or finish
### Background Planes
- `pkg/runtime/bootstrap.go` hard-fails if SecureBus or unified runtime deps are missing.
- `pkg/agent/loop.go` starts `cortex.Cortex` in the background when `AgentLoop.Run` is active.
- `pkg/cortex/cortex.go` schedules memory decay, embedding backfill, consolidation, prune, RL, audit analysis, and drift tasks.
### Subagent Plane
- `pkg/agent/toolloop.go` gives subagents the same unified runtime stack as the parent.
- `pkg/tools/subagent.go` enforces delegation depth and fanout bounds before spawn.
- Subagent completion is published back as a `system` inbound message.
- `pkg/agent/message_router.go` intentionally does **not** forward that system completion to the user; it only logs it.
- Real user-visible subagent communication therefore depends on the subagent using the message/output tools directly.
```mermaid
flowchart TD
CLI["Service.Agent / eval-runner"] --> PD["ProcessDirect"]
GW["Service.Gateway"] --> RUN["go AgentLoop.Run"]
RUN --> BUS["MessageBus inbound"]
BUS --> ROUTER["processMessage"]
PD --> TURN["runAgentLoop"]
ROUTER --> TURN
TURN --> CTX["assembleContext"]
CTX --> BUILD["createFantasyAgent"]
BUILD --> FANTASY["fantasy.Agent.Generate"]
FANTASY --> RT["SecureBus -> Offload -> DAG runtime"]
RT --> TOOLS["Registered tools"]
FANTASY --> POST["postProcess"]
POST --> SAVE["session save"]
POST --> OBS["MaybeObserveAsync"]
POST --> SUM["maybeSummarize"]
POST --> RL["endTask"]
TOOLS --> SUB["SubagentManager"]
SUB --> SYS["system completion message"]
SYS --> ROUTER
```
## 2. ReAct State Machine Truth
- The explicit FSM is defined in `internal/fantasy/react_fsm_machine.go`.
- The state graph is:
- `Init`
- `PrepareStep`
- `LLMCall`
- `ToolValidation`
- `ToolExecution`
- `AppendMessages`
- `StopCheck`
- `Done`
- `Error`
- `internal/fantasy/agent.go` constructs the FSM and emits transitions through optional observers.
- The observer API is real: `WithTransitionObserver`, `WithStepObserver`, and `WithToolResultObserver` exist.
- Production `AgentLoop` does **not** attach any of those observers when building the fantasy agent.
- Result: the FSM exists, the transition log exists, the DB table exists, but runtime transition persistence is not live.
```mermaid
stateDiagram-v2
[*] --> Init
Init --> PrepareStep: Start
PrepareStep --> LLMCall: Prepared
LLMCall --> ToolValidation: LLMResponded
ToolValidation --> ToolExecution: ToolsValidated
ToolExecution --> AppendMessages: ToolsExecuted
AppendMessages --> StopCheck: MessagesAppended
StopCheck --> PrepareStep: Continue
StopCheck --> Done: StopConditionMet / Finished
Init --> Error: Errored
PrepareStep --> Error: Errored
LLMCall --> Error: Errored
ToolValidation --> Error: Errored
ToolExecution --> Error: Errored
AppendMessages --> Error: Errored
StopCheck --> Error: Errored
Error --> Done: Finished
Error --> PrepareStep: RecoveredContinue
```
## 3. Tool Runtime Algorithms and Data Structures
### DAG Tool Execution
- `internal/fantasy/tool_runtime_dag.go` builds a DAG from JSON references like `$tool.<toolCallID>`.
- Dependency resolution is structural, not semantic; it rewrites JSON values, not arbitrary prose.
- Ready nodes are executed in topological waves.
- If a ready node is not parallel-safe, it becomes a barrier and executes alone.
- Parallel-safe ready nodes execute concurrently up to `MaxConcurrency` and then commit in deterministic input order.
- This is effectively:
- dependency analysis
- indegree tracking
- barrier detection
- bounded wave parallelism
- deterministic result commit
### SecureBus
- `pkg/security/securebus/bus.go` is the privilege boundary.
- The synchronous in-process path is:
- capability lookup
- policy validation
- secret injection
- execute underlying tool
- leak scan / redaction
- audit append
- In the current agent runtime, SecureBus is used synchronously through `Bus.Execute(...)`, not via the transport worker queue.
### Offloading
- `pkg/agent/offloading_tool_runtime.go` stores full tool payloads in KV and writes searchable preview rows into `agent_tool_results`.
- Large text results are chunked.
- `tool_result_search` is the live retrieval surface for those stored results.
### Important Runtime Nuance
- `OffloadingToolRuntime` reads the step index from `WithStepIndex(ctx, ...)`.
- `WithStepIndex(...)` is defined, but there are no production call sites.
- Therefore offloaded tool results default to `step_index = 0`.
## 4. Memory System Truth
### Persistence Planes
- Plane 1: session persistence in `pkg/session/manager.go`
- Plane 2: memory tiers in `pkg/memory/store/*` + delegate in `pkg/memory/delegate/sqlite.go`
- Plane 3: runtime/run-state persistence in `pkg/agent/state_store.go`
### Session Persistence
- `SessionManager` keeps in-memory session maps plus optional disk storage plus async DB dual-write.
- Each appended message becomes a `RecallItem`.
- Each appended message is also dual-written as an `ImmutableMessage` when supported.
- Projection pointers are advanced and persisted on append.
- On restore, persisted recall history is replayed chronologically and projection pointers are validated.
### Tiered Memory
- Working context: session-scoped hot memory from delegate-backed working-context records.
- Recall: append-only episodic message records.
- Archival: chunked long-form memory, optionally embedded as `F32_BLOB`.
- Observation: LLM-generated prioritized observations with three-date metadata.
- Knowledge: focus-completion summaries persisted into KV and injected into prompt.
- DAG snapshots: deterministic session compression stored as snapshots, nodes, and edges.
### Observation System
- `pkg/memory/observation/manager.go` is real and asynchronous.
- Observations carry:
- `ObservedAt`
- `ReferencedAt`
- `RelativeDate`
- Observer and reflector are separate components.
- Concurrency is guarded per session by an in-memory running map.
### Deterministic DAG Compression
- `pkg/memory/dag/compress.go` is not LLM-based.
- Algorithm:
- group messages into chunks of 8
- summarize chunk by first sentence per message
- group chunks into sections of 4
- produce optional session summary node
- This is deterministic, reproducible, and auditable.
- Persistent snapshots are hashed and stored via `pkg/memory/dag/store.go`.
### Hybrid Retrieval Router
- `pkg/memory/store/memory_store.go` and `pkg/memory/store/retrieval.go` implement the real retrieval pipeline:
1. keyword search
2. vector search if embedder exists
3. reciprocal rank fusion
4. recency decay
5. metadata filtering
6. hybrid projection search over working context + DAG summaries
7. policy-gated selection between baseline and augmented results
- `pkg/memory/store/retrieval_policy.go` adds a shadow-mode promotion state machine:
- bootstrap in `shadow`
- promote when parity/overlap gates are met
- rollback if promoted quality drops
### Context Tree Reality
- `pkg/contexttree/tree.go` supports:
- semantic + lexical score blending
- time decay
- access-frequency weighting
- type priors
- Boltzmann pruning
- hysteresis
- Production `pkg/agent/summarizer.go` does **not** use the full algorithm.
- Actual production path:
- creates tree nodes with `embedding=nil`
- scores with `queryEmbedding=nil`
- sorts deterministically by score
- greedily fills a token budget
- caches rendered output by `(sessionKey, query, messageCount)`
- So the live system currently uses **lexical + temporal + frequency + type prior** scoring, not semantic scoring.
- Also, live selection does **not** use the package's Boltzmann sampling or hysteresis helpers.
### Context Budget Truth
- `pkg/memory/dag/budget.go` defines a nice budget model for system prompt, observations, knowledge, DAG summaries, raw tail, and tool results.
- `ContextBuilder` separately enforces a system-prompt budget of roughly 40% of context window.
- `applyContextTreeSelection(...)` separately uses `dag.ComputeBudget(...)` to size the raw tail and query-selected history block.
- This means budgeting is currently split across multiple components rather than unified behind one `ActiveContextBuilder`.
## 5. Real Data Structures
- `MessageBus`: buffered inbound/outbound channels plus handler map.
- `SubagentManager`: `tasks` map, `activeChildren` map, next ID counter, depth/fanout caps.
- `ContextTree`: root node, `NodeIndex` map, node children arrays, scoring config.
- `MemoryStore`: delegate, embedder, policy cache, retrieval gates/metrics.
- `DAG`: node map plus root list; persisted as snapshot row + node rows + edge rows.
- `SessionManager`: in-memory session map, LRU, async persist queue, projection-pointer records.
- `StateStore`: `agent_runs`, `agent_run_states`, `agent_state_transitions`, `agent_tool_results`, `agent_checkpoints`.
- `SecureBus`: policy engine, secret store, redactor, audit log, optional transport.
- `ProjectionContract`: `ActiveContextProjection`, `ProjectionSegment`, `ImmutableSpanRef` in `pkg/memory/kernel_contract.go`.
## 6. Reconciliation Matrix
### Shipped and true
- Unified kernel runtime with fail-fast boot invariants.
- SecureBus mediation for tool execution.
- Offloaded tool result persistence and retrieval.
- DAG tool execution with dependency-aware parallel waves.
- Observational memory with async manager.
- Focus primitives with knowledge-block accumulation.
- Skill graph tooling: `skill_search`, `skill_read`, `skill_traverse`.
- Retrieval tools: `keyword_search`, `semantic_search`, `chunk_read`.
- Persistent DAG snapshots and DAG tools: `dag_expand`, `dag_describe`, `dag_grep`.
- Projection pointers and deterministic session restore validation.
- Subagent depth/fanout/runtime parity guardrails.
- Promptfoo harness with `maxConcurrency: 1` in `eval/promptfooconfig.yaml`.
### Present but partial
- Run-state persistence is live, but the step indexing story is incomplete.
- ContextTree package is richer than its actual production use.
- Kernel projection contracts exist as types, not as a production builder.
- Checkpoint and conversation fork/merge scaffolding exist, but runtime checkpoint creation is not wired.
- Eval hardening plan exists, but docs do not clearly reconcile plan-vs-current behavior.
### Present in docs more than code
- RLM integration.
- some ADR-001 convergence claims around live RLM-backed execution.
- README migration count and config sample.
- roadmap sections that still read as future work even where tools/features are already shipped.
### Present in code but not wired
- `WithTransitionObserver` / `AddTransition`.
- `WithStepIndex(...)`.
- `ActiveContextBuilder`.
- `CheckpointStore` production usage.
- `pkg/rlm` production imports outside tests.
## 7. Counterexample-Style Findings
### Finding 1: `TransitionPersistenceComplete` is false
- Claim: every FSM transition is persisted.
- Counterexample: `internal/fantasy/agent.go` supports transition observers, and `pkg/agent/state_store.go` supports `AddTransition`, but `pkg/agent/agent_run.go` never attaches `WithTransitionObserver(...)`.
- Effect: `agent_state_transitions` is effectively dead schema in production.
### Finding 2: `MonotonicStepIndex` is false
- Claim: persisted run/tool records preserve the true ReAct step number.
- Counterexample: `pkg/agent/offloading_tool_runtime.go` reads step index from context via `StepIndexFromCtx(...)`, but `WithStepIndex(...)` has no production callers.
- Counterexample: `pkg/agent/securebus_runtime.go` has a `StepIndex` field, but the fantasy agent is created with no per-step updates to that field.
- Effect: `agent_tool_results.step_index` is effectively always `0`, and `agent_run_states.step_index` is not a trustworthy global turn-step index.
### Finding 3: `SemanticContextSelection` is overstated
- Claim: ContextTree production selection is semantic + lexical.
- Counterexample: `pkg/agent/summarizer.go` passes `nil` embeddings into `ContextTree`.
- Effect: live selection is lexical/temporal/frequency/type-prior only.
### Finding 4: `ContextBudgetingIsUnified` is false
- Claim: one controller owns active-context assembly.
- Counterexample: `ContextBuilder`, `applyContextTreeSelection(...)`, DAG budget helpers, and working-context injection each own part of the budget story.
- Effect: the architecture is effective but still transitional, not yet a single controller model.
### Finding 5: `RLMIntegrated` is false
- Claim: production agent runtime uses `pkg/rlm`.
- Counterexample: repo-wide imports of `pkg/rlm` outside tests are absent.
- Effect: RLM is a tested subsystem and an architectural direction, not current runtime truth.
### Finding 6: `CheckpointableRuntime` is partial
- Claim: runtime can be checkpointed and forked as a normal live feature.
- Counterexample: checkpoint APIs and fork/merge storage exist, but there is no production path that creates checkpoints during normal agent execution.
- Effect: the storage layer exists ahead of the runtime behavior.
### Finding 7: `RLToolCallMetric` is imprecise
- `pkg/agent/agent_run.go` records `TaskCompletion.ToolCalls = stepCount`.
- Effect: RL/task analytics are counting ReAct steps, not true tool invocations.
## 8. Documentation Drift
- `README.md` still says Goose has 10 versioned migrations; the repo now has 16 numbered migrations plus context wiring.
- `README.md` includes `"tools.progressive_disclosure": true`, but `pkg/config/config.go` has no such field. Progressive disclosure is effectively runtime behavior, not a config toggle.
- `README.md` and ADR-001 imply RLM is part of the live execution story; code does not support that.
- `ROADMAP.md` still treats some already-shipped features as future work and leaves others underspecified relative to code truth.
- `docs/adr/001-isolated-tool-runtime.md` is still `Proposed` even though large parts of the secure runtime architecture are already live.
- `eval/README.md` does not fully reconcile with the signal-first hardening plan that discussed `35/38`, generic fallback leakage, and concurrency/determinism concerns.
- The hardening plan references `eval/promptfooconfig-default.yaml`, but the live repo uses `eval/promptfooconfig.yaml`.
## 9. TLA+ Spec Brief
### System
- A single agent runtime that receives messages, assembles bounded context, executes zero or more tool steps through a mediated runtime, persists turn artifacts, and may delegate to subagents.
### Actors
- User / external channel
- `AgentLoop`
- `fantasy.Agent` ReAct controller
- `SecureBus`
- tool runtime / tools
- `SessionManager`
- `MemoryStore`
- `SubagentManager`
- `Cortex`
- environment failures: DB errors, model errors, tool errors, cancellation
### State
- inbound queue
- outbound queue
- sessions and immutable history
- working context / observations / knowledge / DAG snapshots
- active runs and run states
- persisted tool results
- active subagent tasks
- audit log
### Safety invariants worth checking
- `I_SecureBusMediatesAllToolExec`
- `I_DelegationDepthBound`
- `I_DelegationFanoutBound`
- `I_ProjectionPointerMonotonic`
- `I_ToolResultRowHasReachableFullKey`
- `I_NoOrphanToolMessageAfterTruncation`
- `I_TransitionPersistenceComplete` currently false
- `I_MonotonicStepIndex` currently false
### Liveness properties worth checking
- If a user message is accepted and the model/tool runtime eventually returns, the run eventually reaches `Done` or explicit `Error`.
- If a subagent is spawned and its run loop terminates, the parent eventually receives a completion announcement.
- If observation threshold is crossed and storage/model calls succeed, observation state eventually updates.
- If hard compaction threshold is crossed and summarization succeeds within bounded cycles, token pressure eventually drops below the critical threshold.
### Minimal TLC model bounds
- 1 session
- at most 3 ReAct steps
- at most 2 tool calls per step
- at most 1 subagent child per step
- max delegation depth 2
- max DAG ready set 2
- bounded observation list of 4 items
- bounded history of 6 messages
```tla
---- MODULE DragonScaleKernel ----
EXTENDS Sequences, FiniteSets, TLC
VARIABLES inboundQ, sessions, runs, runStates, toolResults, subagents, audits
Init ==
/\ inboundQ = << >>
/\ sessions = [s \in {"s0"} |-> << >>]
/\ runs = [r \in {} |-> [status |-> "none"]]
/\ runStates = {}
/\ toolResults = {}
/\ subagents = {}
/\ audits = << >>
ReceiveUser(msg) == inboundQ' = Append(inboundQ, msg)
StartTurn == \E msg \in SeqToSet(inboundQ): TRUE
ExecuteTool == TRUE
PersistToolResult == TRUE
CompleteTurn == TRUE
SpawnSubagent == TRUE
FinishSubagent == TRUE
Next ==
\/ \E msg : ReceiveUser(msg)
\/ StartTurn
\/ ExecuteTool
\/ PersistToolResult
\/ CompleteTurn
\/ SpawnSubagent
\/ FinishSubagent
====
```
## 10. Sequential Priority Order
1. Fix runtime observability first.
- Wire `WithTransitionObserver(...)` from `pkg/agent/agent_run.go` into `StateStore.AddTransition(...)`.
- Propagate true step numbers into `SecureBusToolRuntime` and `OffloadingToolRuntime`.
- Change RL `ToolCalls` accounting from step count to actual tool result count.
2. Reconcile the context model second.
- Decide whether `ActiveContextBuilder` becomes real or gets downgraded to a future contract.
- Decide whether ContextTree remains lexical-only or receives real embeddings.
- Decide whether DAG prompt rendering should be direct runtime input or remain a search/recovery structure.
3. Reconcile docs third.
- Update `README.md`, `ROADMAP.md`, `ADR-001`, and `eval/README.md` to code-truth.
- Explicitly mark RLM as “present in package, not wired in production” unless that changes.
4. Decide the fate of RLM and checkpoints fourth.
- Either wire `pkg/rlm` and runtime checkpoints for real, or demote those claims out of the production architecture story.
5. Refresh live truth last.
- Re-run `go test ./...`
- Re-run promptfoo
- Re-check any external tracker state only after code/docs truth are aligned
## Bottom Line
- DragonScale already has a real, non-trivial kernel: mediated tool execution, async memory maintenance, persistent session continuity, retrieval routing, focus/obligation/skill graph tooling, and subagent guardrails are all live.
- The repo is **not** yet in a single perfectly reconciled architectural state.
- The main unresolved gap is that the runtime's **bookkeeping and docs lag the actual kernel**: transitions are not persisted, step indexes are not authoritative, RLM is not wired, and documentation still blends shipped behavior with intended behavior.

View file

@ -12,6 +12,8 @@ npm install -g promptfoo
make eval make eval
``` ```
`make eval` auto-builds `eval/bin/eval-runner` when it is missing.
To show richer promptfoo output with progress bars (when supported), run: To show richer promptfoo output with progress bars (when supported), run:
```bash ```bash
@ -125,8 +127,12 @@ Create a new YAML file in `eval/cases/` following this pattern:
- type: javascript - type: javascript
value: | value: |
const trace = JSON.parse(output); const trace = JSON.parse(output);
// Return { pass: bool, score: 0-1, reason: string } const usedTool = trace.metrics.tool_call_count > 0;
return { pass: true, score: 1.0, reason: 'explanation' }; return {
pass: usedTool,
score: usedTool ? 1.0 : 0.0,
reason: usedTool ? 'tool call observed' : 'expected at least one tool call'
};
``` ```
For generated long-context suites: For generated long-context suites:
@ -137,7 +143,9 @@ python eval/scripts/generate_long_context_cases.py --count 12 --seed 20260221
## A/B Comparison ## A/B Comparison
`make eval-compare` builds both your current branch and main, then runs the identical test suite against both. Results show a side-by-side comparison matrix with per-test scores. `make eval-compare` builds both your current branch and `main`, then runs the identical test suite against both. Results show a side-by-side comparison matrix with per-test scores.
The comparison flow uses a temporary git worktree for `main`, so the active checkout is never stashed or branch-switched underneath the user.
## Environment Variables ## Environment Variables

View file

@ -26,7 +26,7 @@
}; };
const toolNames = toolCalls.map(getToolName); const toolNames = toolCalls.map(getToolName);
const usedMemoryLikeTool = toolNames.includes('memory') || toolNames.includes('keyword_search') || toolNames.includes('semantic_search'); const usedMemoryLikeTool = toolNames.includes('memory') || toolNames.includes('keyword_search') || toolNames.includes('semantic_search');
return { pass: true, score: usedMemoryLikeTool ? 1.0 : 0.5, reason: usedMemoryLikeTool ? 'used memory-capable tooling for commitments' : `no memory tool observed (tools: ${toolNames.join(', ')})` }; return { pass: usedMemoryLikeTool, score: usedMemoryLikeTool ? 1.0 : 0.0, reason: usedMemoryLikeTool ? 'used memory-capable tooling for commitments' : `no memory tool observed (tools: ${toolNames.join(', ')})` };
- description: "follow-up escalation after missed commitment" - description: "follow-up escalation after missed commitment"
vars: vars:

View file

@ -45,10 +45,10 @@
const trace = JSON.parse(output); const trace = JSON.parse(output);
const out = (trace.output || '').toLowerCase(); const out = (trace.output || '').toLowerCase();
const dur = trace.metrics.total_duration_ms; const dur = trace.metrics.total_duration_ms;
const notHung = dur < 130000; const notHung = dur < 60000;
const graceful = out.includes('timeout') || out.includes('cancel') || out.includes('too long') || const graceful = out.includes('timeout') || out.includes('cancel') || out.includes('too long') ||
out.includes('killed') || out.includes('error') || out.includes('interrupt') || out.includes('killed') || out.includes('error') || out.includes('interrupt') ||
out.length > 5; out.includes('denied') || out.includes('not permitted') || out.includes('not allowed');
return { pass: notHung && graceful, score: notHung ? 1.0 : 0.0, reason: `duration=${dur}ms, not_hung=${notHung}, response=${graceful}` }; return { pass: notHung && graceful, score: notHung ? 1.0 : 0.0, reason: `duration=${dur}ms, not_hung=${notHung}, response=${graceful}` };
- description: "invalid tool args: schema validation rejection" - description: "invalid tool args: schema validation rejection"

View file

@ -45,7 +45,9 @@
if (isDefault) return { pass: false, score: 0.2, reason: 'fell back to default empty response' }; if (isDefault) return { pass: false, score: 0.2, reason: 'fell back to default empty response' };
const graceful = out.includes('no results') || out.includes('nothing') || out.includes('not found') || const graceful = out.includes('no results') || out.includes('nothing') || out.includes('not found') ||
out.includes("didn't find") || out.includes("don't have") || out.includes('no memories') || out.includes("didn't find") || out.includes("don't have") || out.includes('no memories') ||
out.includes('no matching') || out.length > 10; out.includes('no matching') || out.includes('no stored memories') ||
out.includes('no stored memory') || out.includes('no entries') ||
out.includes('no stored information') || out.includes('no actual stored information');
return { pass: graceful, score: graceful ? 1.0 : 0.0, reason: graceful ? 'handled empty search gracefully' : 'no meaningful response' }; return { pass: graceful, score: graceful ? 1.0 : 0.0, reason: graceful ? 'handled empty search gracefully' : 'no meaningful response' };
- description: "agent responds to greeting with memory system active" - description: "agent responds to greeting with memory system active"

View file

@ -19,6 +19,7 @@ import (
func main() { func main() {
logger.SetLevel(logger.ERROR) logger.SetLevel(logger.ERROR)
_ = os.Setenv("DRAGONSCALE_EVAL_RUNTIME", "1")
prompt, err := resolvePrompt() prompt, err := resolvePrompt()
if err != nil { if err != nil {
@ -51,6 +52,7 @@ func emptyPromptTrace(prompt string) *instrumentation.Trace {
} }
return &instrumentation.Trace{ return &instrumentation.Trace{
Output: "No prompt provided. Please provide a message.", Output: "No prompt provided. Please provide a message.",
Steps: []instrumentation.TraceStep{},
Metrics: instrumentation.Metrics{ Metrics: instrumentation.Metrics{
TotalDurationMs: 0, TotalDurationMs: 0,
}, },
@ -58,7 +60,24 @@ func emptyPromptTrace(prompt string) *instrumentation.Trace {
} }
func resolveEvalConfig() (*config.Config, error) { func resolveEvalConfig() (*config.Config, error) {
return dragonruntime.LoadEvalConfig(evalRunnerTimeout()) cfg, err := dragonruntime.LoadEvalConfig(evalRunnerTimeout())
if err != nil {
return nil, err
}
stabilizeEvalConfig(cfg)
return cfg, nil
}
func stabilizeEvalConfig(cfg *config.Config) {
if cfg == nil {
return
}
if cfg.Agents.Defaults.Temperature != 0 {
cfg.Agents.Defaults.Temperature = 0
}
if cfg.Agents.Defaults.MaxToolIterations <= 0 || cfg.Agents.Defaults.MaxToolIterations > 8 {
cfg.Agents.Defaults.MaxToolIterations = 8
}
} }
func readPrompt() (string, error) { func readPrompt() (string, error) {

View file

@ -0,0 +1,40 @@
package main
import (
"encoding/json"
"strings"
"testing"
"github.com/ZanzyTHEbar/dragonscale/pkg/config"
)
func TestEmptyPromptTraceEmitsEmptyStepsArray(t *testing.T) {
trace := emptyPromptTrace(" ")
if trace == nil {
t.Fatal("expected empty prompt trace")
}
raw, err := json.Marshal(trace)
if err != nil {
t.Fatalf("marshal trace: %v", err)
}
if !strings.Contains(string(raw), `"steps":[]`) {
t.Fatalf("expected empty steps array, got %s", string(raw))
}
}
func TestStabilizeEvalConfigClampsAgentDefaults(t *testing.T) {
cfg := config.DefaultConfig()
cfg.Agents.Defaults.Temperature = 0.7
cfg.Agents.Defaults.MaxToolIterations = 20
stabilizeEvalConfig(cfg)
if cfg.Agents.Defaults.Temperature != 0 {
t.Fatalf("expected eval temperature 0, got %v", cfg.Agents.Defaults.Temperature)
}
if cfg.Agents.Defaults.MaxToolIterations != 8 {
t.Fatalf("expected max tool iterations 8, got %d", cfg.Agents.Defaults.MaxToolIterations)
}
}

View file

@ -21,7 +21,7 @@ defaultTest:
value: | value: |
try { try {
const trace = JSON.parse(output); const trace = JSON.parse(output);
const valid = trace.hasOwnProperty('output') && trace.hasOwnProperty('metrics'); const valid = typeof trace.output === 'string' && Array.isArray(trace.steps) && trace.metrics && typeof trace.metrics.total_duration_ms === 'number';
return { pass: valid, score: valid ? 1.0 : 0.0, reason: valid ? 'valid trace JSON' : 'invalid trace structure' }; return { pass: valid, score: valid ? 1.0 : 0.0, reason: valid ? 'valid trace JSON' : 'invalid trace structure' };
} catch(e) { } catch(e) {
return { pass: false, score: 0, reason: 'output is not valid JSON: ' + e.message }; return { pass: false, score: 0, reason: 'output is not valid JSON: ' + e.message };
@ -33,7 +33,8 @@ defaultTest:
const trace = JSON.parse(output); const trace = JSON.parse(output);
const dur = trace.metrics.total_duration_ms; const dur = trace.metrics.total_duration_ms;
const score = dur < 30000 ? 1.0 : dur < 90000 ? 1.0 - (dur - 30000) / 60000 : 0.0; const score = dur < 30000 ? 1.0 : dur < 90000 ? 1.0 - (dur - 30000) / 60000 : 0.0;
return { pass: true, score, reason: `duration: ${dur}ms (score: ${score.toFixed(2)})` }; const pass = dur < 90000;
return { pass, score, reason: `duration: ${dur}ms (score: ${score.toFixed(2)})` };
tests: "cases/*.yaml" tests: "cases/*.yaml"

View file

@ -0,0 +1,81 @@
description: "DragonScale targeted re-verification"
maxConcurrency: 1
providers:
- id: "exec:./bin/eval-runner"
label: "dragonscale"
config:
timeout: 180000
env:
DRAGONSCALE_EVAL_CONFIG: "./configs/default.json"
defaultTest:
assert:
- type: javascript
value: |
try {
const trace = JSON.parse(output);
const valid = typeof trace.output === 'string' && Array.isArray(trace.steps) && trace.metrics && typeof trace.metrics.total_duration_ms === 'number';
return { pass: valid, score: valid ? 1.0 : 0.0, reason: valid ? 'valid trace JSON' : 'invalid trace structure' };
} catch(e) {
return { pass: false, score: 0, reason: 'output is not valid JSON: ' + e.message };
}
- type: javascript
value: |
const trace = JSON.parse(output);
const dur = trace.metrics.total_duration_ms;
const score = dur < 30000 ? 1.0 : dur < 90000 ? 1.0 - (dur - 30000) / 60000 : 0.0;
const pass = dur < 90000;
return { pass, score, reason: `duration: ${dur}ms (score: ${score.toFixed(2)})` };
tests:
- description: "edit file: agent edits an existing file"
vars:
prompt: "First write a file called edit_target.txt with 'hello world'. Then edit it to replace 'world' with 'dragonscale'. Read it back and confirm."
assert:
- type: javascript
value: |
const trace = JSON.parse(output);
if (trace.error) return { pass: false, score: 0, reason: trace.error };
const toolCalls = trace.steps.filter(s => s.type === 'tool_call');
const getToolName = (t) => {
if (t.tool !== 'tool_call') return t.tool;
try { const a = typeof t.args === 'string' ? JSON.parse(t.args) : t.args; return (a && a.tool_name) || t.tool; } catch(e) { return t.tool; }
};
const toolNames = toolCalls.map(getToolName);
const hasEdit = toolNames.includes('edit_file');
const hasWrite = toolNames.includes('write_file');
const hasAnyModify = hasEdit || hasWrite;
return { pass: hasAnyModify, score: hasEdit ? 1.0 : (hasWrite ? 0.7 : 0.0), reason: `edit=${hasEdit}, write=${hasWrite} (tools: ${toolNames.join(', ')})` };
- type: javascript
value: |
const trace = JSON.parse(output);
const out = (trace.output || '').toLowerCase();
const hasDragonScale = out.includes('dragonscale');
return { pass: hasDragonScale, score: hasDragonScale ? 1.0 : 0.0, reason: hasDragonScale ? 'confirmed edit result' : 'did not confirm dragonscale in output' };
- description: "skill read: load skill content"
vars:
prompt: "Read the 'eval-test-skill' skill and tell me what greeting templates it provides."
assert:
- type: javascript
value: |
const trace = JSON.parse(output);
if (trace.error) return { pass: false, score: 0, reason: trace.error };
const toolCalls = trace.steps.filter(s => s.type === 'tool_call');
const getToolName = (t) => {
if (t.tool !== 'tool_call') return t.tool;
try { const a = typeof t.args === 'string' ? JSON.parse(t.args) : t.args; return (a && a.tool_name) || t.tool; } catch(e) { return t.tool; }
};
const toolNames = toolCalls.map(getToolName);
const usedSkillRead = toolNames.includes('skill_read');
return { pass: usedSkillRead, score: usedSkillRead ? 1.0 : 0.0, reason: usedSkillRead ? 'used skill_read' : `no skill_read (tools: ${toolNames.join(', ')})` };
- type: javascript
value: |
const trace = JSON.parse(output);
const out = (trace.output || '').toLowerCase();
const hasContent = out.includes('formal') || out.includes('casual') || out.includes('greeting') || out.includes('template');
return { pass: hasContent, score: hasContent ? 1.0 : 0.0, reason: hasContent ? 'returned skill content' : 'did not return skill content' };
outputPath: "results/reverify-two-cases.json"

View file

@ -10,11 +10,17 @@ PROJECT_ROOT="$(dirname "$EVAL_DIR")"
REPEAT=${1:-3} REPEAT=${1:-3}
NPM_CMD="${EVAL_NPM_CMD:-npx}" NPM_CMD="${EVAL_NPM_CMD:-npx}"
read -r -a NPM_CMD_ARR <<< "${NPM_CMD}" read -r -a NPM_CMD_ARR <<< "${NPM_CMD}"
PROMPTFOO_ARGS="${DRAGONSCALE_PROMPTFOO_ARGS:---no-cache --no-progress-bar -j 1}"
read -r -a PROMPTFOO_ARGS_ARR <<< "${PROMPTFOO_ARGS}"
export DEVCONTAINER_EXEC="" export DEVCONTAINER_EXEC=""
TEMP_CONFIG="$(mktemp "${SCRIPT_DIR}/promptfoo-compare-XXXXXX.yaml")" TEMP_CONFIG="$(mktemp "${SCRIPT_DIR}/promptfoo-compare-XXXXXX.yaml")"
TEMP_WORKTREE=""
cleanup_compare_config() { cleanup_compare_config() {
rm -f "$TEMP_CONFIG" rm -f "$TEMP_CONFIG"
if [[ -n "$TEMP_WORKTREE" && -d "$TEMP_WORKTREE" ]]; then
git -C "$PROJECT_ROOT" worktree remove --force "$TEMP_WORKTREE" >/dev/null 2>&1 || rm -rf "$TEMP_WORKTREE"
fi
} }
trap cleanup_compare_config EXIT INT TERM trap cleanup_compare_config EXIT INT TERM
@ -33,20 +39,12 @@ echo "[1/4] Building eval-runner from current branch..."
make DEVCONTAINER_EXEC= eval-build 2>&1 | tail -1 make DEVCONTAINER_EXEC= eval-build 2>&1 | tail -1
cp "$EVAL_DIR/bin/eval-runner" "$EVAL_DIR/bin/eval-runner-branch" cp "$EVAL_DIR/bin/eval-runner" "$EVAL_DIR/bin/eval-runner-branch"
# 2. Build main branch eval-runner # 2. Build main branch eval-runner in an isolated worktree
CURRENT_BRANCH=$(git branch --show-current)
STASH_RESULT=$(git stash 2>&1)
echo "[2/4] Building eval-runner from main branch..." echo "[2/4] Building eval-runner from main branch..."
git checkout main 2>/dev/null TEMP_WORKTREE="$(mktemp -d "${TMPDIR:-/tmp}/dragonscale-eval-main-XXXXXX")"
make DEVCONTAINER_EXEC= eval-build 2>&1 | tail -1 git -C "$PROJECT_ROOT" worktree add --force --detach "$TEMP_WORKTREE" main >/dev/null
cp "$EVAL_DIR/bin/eval-runner" "$EVAL_DIR/bin/eval-runner-main" make -C "$TEMP_WORKTREE" DEVCONTAINER_EXEC= eval-build 2>&1 | tail -1
cp "$TEMP_WORKTREE/eval/bin/eval-runner" "$EVAL_DIR/bin/eval-runner-main"
# Restore working branch
git checkout "$CURRENT_BRANCH" 2>/dev/null
if [[ "$STASH_RESULT" != "No local changes to save" ]]; then
git stash pop 2>/dev/null || true
fi
# Put branch binary back as the default eval-runner # Put branch binary back as the default eval-runner
cp "$EVAL_DIR/bin/eval-runner-branch" "$EVAL_DIR/bin/eval-runner" cp "$EVAL_DIR/bin/eval-runner-branch" "$EVAL_DIR/bin/eval-runner"
@ -69,6 +67,8 @@ fi
cat > "$TEMP_CONFIG" <<YAML cat > "$TEMP_CONFIG" <<YAML
description: "DragonScale A/B comparison (branch vs main)" description: "DragonScale A/B comparison (branch vs main)"
maxConcurrency: 1
providers: providers:
- id: "exec:./bin/eval-runner" - id: "exec:./bin/eval-runner"
label: "branch" label: "branch"
@ -91,7 +91,7 @@ defaultTest:
value: | value: |
try { try {
const trace = JSON.parse(output); const trace = JSON.parse(output);
const valid = trace.hasOwnProperty('output') && trace.hasOwnProperty('metrics'); const valid = typeof trace.output === 'string' && Array.isArray(trace.steps) && trace.metrics && typeof trace.metrics.total_duration_ms === 'number';
return { pass: valid, score: valid ? 1.0 : 0.0, reason: valid ? 'valid trace' : 'invalid trace' }; return { pass: valid, score: valid ? 1.0 : 0.0, reason: valid ? 'valid trace' : 'invalid trace' };
} catch(e) { } catch(e) {
return { pass: false, score: 0, reason: 'not JSON: ' + e.message }; return { pass: false, score: 0, reason: 'not JSON: ' + e.message };
@ -100,8 +100,9 @@ defaultTest:
value: | value: |
const trace = JSON.parse(output); const trace = JSON.parse(output);
const dur = trace.metrics.total_duration_ms; const dur = trace.metrics.total_duration_ms;
const ok = dur < 60000; const score = dur < 30000 ? 1.0 : dur < 90000 ? 1.0 - (dur - 30000) / 60000 : 0.0;
return { pass: ok, score: ok ? 1.0 : 0.0, reason: `${dur}ms` }; const pass = dur < 90000;
return { pass, score, reason: `${dur}ms (score: ${score.toFixed(2)})` };
transform: "JSON.stringify({ prompt: vars.prompt })" transform: "JSON.stringify({ prompt: vars.prompt })"
tests: "cases/*.yaml" tests: "cases/*.yaml"
@ -110,7 +111,7 @@ YAML
# 4. Run comparison # 4. Run comparison
echo "[3/4] Running eval comparison (${REPEAT}x)..." echo "[3/4] Running eval comparison (${REPEAT}x)..."
"${NPM_CMD_ARR[@]}" promptfoo eval -c "$TEMP_CONFIG" --repeat "$REPEAT" --no-progress-bar "${NPM_CMD_ARR[@]}" promptfoo eval -c "$TEMP_CONFIG" --repeat "$REPEAT" "${PROMPTFOO_ARGS_ARR[@]}"
echo "" echo ""
echo "[4/4] Results saved to eval/results/comparison.json" echo "[4/4] Results saved to eval/results/comparison.json"

View file

@ -5,11 +5,12 @@ import (
"context" "context"
"errors" "errors"
"fmt" "fmt"
jsonv2 "github.com/go-json-experiment/json"
"maps" "maps"
"slices" "slices"
"sync" "sync"
jsonv2 "github.com/go-json-experiment/json"
"charm.land/fantasy/schema" "charm.land/fantasy/schema"
) )
@ -392,6 +393,7 @@ func (a *agent) Generate(ctx context.Context, opts AgentCall) (*AgentResult, err
fsm.Fire(ctx, ReActTriggerStart) fsm.Fire(ctx, ReActTriggerStart)
for { for {
stepIdx = len(steps)
stepInputMessages := append(initialPrompt, responseMessages...) stepInputMessages := append(initialPrompt, responseMessages...)
stepModel := a.settings.model stepModel := a.settings.model
stepSystemPrompt := a.settings.systemPrompt stepSystemPrompt := a.settings.systemPrompt
@ -493,10 +495,11 @@ func (a *agent) Generate(ctx context.Context, opts AgentCall) (*AgentResult, err
fsm.Fire(ctx, ReActTriggerToolsValidated) fsm.Fire(ctx, ReActTriggerToolsValidated)
var toolResults []ToolResultContent var toolResults []ToolResultContent
stepCtx := WithStepIndex(ctx, len(steps))
if a.settings.toolRuntime != nil { if a.settings.toolRuntime != nil {
toolResults, err = a.settings.toolRuntime.Execute(ctx, stepTools, stepToolCalls, nil) toolResults, err = a.settings.toolRuntime.Execute(stepCtx, stepTools, stepToolCalls, nil)
} else { } else {
toolResults, err = a.executeTools(ctx, stepTools, stepToolCalls, nil) toolResults, err = a.executeTools(stepCtx, stepTools, stepToolCalls, nil)
} }
fsm.Fire(ctx, ReActTriggerToolsExecuted) fsm.Fire(ctx, ReActTriggerToolsExecuted)
@ -536,7 +539,6 @@ func (a *agent) Generate(ctx context.Context, opts AgentCall) (*AgentResult, err
Messages: currentStepMessages, Messages: currentStepMessages,
} }
steps = append(steps, stepResult) steps = append(steps, stepResult)
stepIdx = len(steps) - 1
for _, obs := range a.settings.stepObservers { for _, obs := range a.settings.stepObservers {
obs.OnReActStep(ctx, len(steps)-1, stepResult) obs.OnReActStep(ctx, len(steps)-1, stepResult)
@ -816,6 +818,7 @@ func (a *agent) Stream(ctx context.Context, opts AgentStreamCall) (*AgentResult,
} }
for stepNumber := 0; ; stepNumber++ { for stepNumber := 0; ; stepNumber++ {
streamStepIdx = len(steps)
stepInputMessages := append(initialPrompt, responseMessages...) stepInputMessages := append(initialPrompt, responseMessages...)
stepModel := a.settings.model stepModel := a.settings.model
stepSystemPrompt := a.settings.systemPrompt stepSystemPrompt := a.settings.systemPrompt
@ -931,7 +934,6 @@ func (a *agent) Stream(ctx context.Context, opts AgentStreamCall) (*AgentResult,
streamFSM.Fire(ctx, ReActTriggerToolsExecuted) streamFSM.Fire(ctx, ReActTriggerToolsExecuted)
steps = append(steps, result.StepResult) steps = append(steps, result.StepResult)
streamStepIdx = len(steps) - 1
totalUsage = addUsage(totalUsage, result.StepResult.Usage) totalUsage = addUsage(totalUsage, result.StepResult.Usage)
for _, obs := range a.settings.stepObservers { for _, obs := range a.settings.stepObservers {
@ -1229,6 +1231,7 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op
activeReasoningContent := make(map[string]reasoningContent) activeReasoningContent := make(map[string]reasoningContent)
useRuntimeBatch := a.settings.toolRuntime != nil useRuntimeBatch := a.settings.toolRuntime != nil
stepCtx := WithStepIndex(ctx, len(steps))
// Set up concurrent tool execution (used only when no ToolRuntime is set) // Set up concurrent tool execution (used only when no ToolRuntime is set)
type toolExecutionRequest struct { type toolExecutionRequest struct {
@ -1261,7 +1264,7 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op
parallelSem <- struct{}{} parallelSem <- struct{}{}
toolExecutionWg.Go(func() { toolExecutionWg.Go(func() {
defer func() { <-parallelSem }() defer func() { <-parallelSem }()
result, isCriticalError := a.executeSingleTool(ctx, toolMap, req.toolCall, opts.OnToolResult) result, isCriticalError := a.executeSingleTool(stepCtx, toolMap, req.toolCall, opts.OnToolResult)
toolStateMu.Lock() toolStateMu.Lock()
toolResults = append(toolResults, result) toolResults = append(toolResults, result)
if isCriticalError && toolExecutionErr == nil { if isCriticalError && toolExecutionErr == nil {
@ -1273,7 +1276,7 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op
}) })
} else { } else {
sequentialMu.Lock() sequentialMu.Lock()
result, isCriticalError := a.executeSingleTool(ctx, toolMap, req.toolCall, opts.OnToolResult) result, isCriticalError := a.executeSingleTool(stepCtx, toolMap, req.toolCall, opts.OnToolResult)
toolStateMu.Lock() toolStateMu.Lock()
toolResults = append(toolResults, result) toolResults = append(toolResults, result)
if isCriticalError && toolExecutionErr == nil { if isCriticalError && toolExecutionErr == nil {
@ -1487,7 +1490,7 @@ func (a *agent) processStepStream(ctx context.Context, stream StreamResponse, op
if useRuntimeBatch { if useRuntimeBatch {
if len(stepToolCalls) > 0 { if len(stepToolCalls) > 0 {
var err error var err error
toolResults, err = a.settings.toolRuntime.Execute(ctx, stepTools, stepToolCalls, opts.OnToolResult) toolResults, err = a.settings.toolRuntime.Execute(stepCtx, stepTools, stepToolCalls, opts.OnToolResult)
if err != nil { if err != nil {
return stepExecutionResult{}, err return stepExecutionResult{}, err
} }

View file

@ -0,0 +1,173 @@
package fantasy
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type recordingToolRuntime struct {
stepIndices []int
}
func (r *recordingToolRuntime) Execute(ctx context.Context, _ []AgentTool, calls []ToolCallContent, _ func(ToolResultContent) error) ([]ToolResultContent, error) {
if len(calls) == 0 {
return nil, nil
}
r.stepIndices = append(r.stepIndices, StepIndexFromCtx(ctx))
results := make([]ToolResultContent, 0, len(calls))
for _, call := range calls {
results = append(results, ToolResultContent{
ToolCallID: call.ToolCallID,
ToolName: call.ToolName,
Result: ToolResultOutputContentText{Text: "ok:" + call.ToolCallID},
})
}
return results, nil
}
type multiStepToolModel struct{}
func (m *multiStepToolModel) Generate(_ context.Context, call Call) (*Response, error) {
toolResults := countToolResults(call.Prompt)
switch toolResults {
case 0:
return toolCallResponse("call-1"), nil
case 1:
return toolCallResponse("call-2"), nil
default:
return &Response{
Content: ResponseContent{TextContent{Text: "final answer"}},
FinishReason: FinishReasonStop,
}, nil
}
}
func (m *multiStepToolModel) Stream(ctx context.Context, call Call) (StreamResponse, error) {
resp, err := m.Generate(ctx, call)
if err != nil {
return nil, err
}
return func(yield func(StreamPart) bool) {
if len(resp.Content.ToolCalls()) > 0 {
for _, tc := range resp.Content.ToolCalls() {
if !yield(StreamPart{
Type: StreamPartTypeToolCall,
ID: tc.ToolCallID,
ToolCallName: tc.ToolName,
ToolCallInput: tc.Input,
}) {
return
}
}
yield(StreamPart{Type: StreamPartTypeFinish, FinishReason: FinishReasonToolCalls})
return
}
text := resp.Content.Text()
if !yield(StreamPart{Type: StreamPartTypeTextStart, ID: "text-0"}) {
return
}
if !yield(StreamPart{Type: StreamPartTypeTextDelta, ID: "text-0", Delta: text}) {
return
}
if !yield(StreamPart{Type: StreamPartTypeTextEnd, ID: "text-0"}) {
return
}
yield(StreamPart{Type: StreamPartTypeFinish, FinishReason: FinishReasonStop})
}, nil
}
func (m *multiStepToolModel) GenerateObject(_ context.Context, _ ObjectCall) (*ObjectResponse, error) {
return nil, nil
}
func (m *multiStepToolModel) StreamObject(_ context.Context, _ ObjectCall) (ObjectStreamResponse, error) {
return nil, nil
}
func (m *multiStepToolModel) Provider() string { return "mock" }
func (m *multiStepToolModel) Model() string { return "multi-step-tool-model" }
func toolCallResponse(id string) *Response {
return &Response{
Content: ResponseContent{
ToolCallContent{
ToolCallID: id,
ToolName: "echo",
Input: `{}`,
},
},
FinishReason: FinishReasonToolCalls,
}
}
func countToolResults(prompt []Message) int {
count := 0
for _, msg := range prompt {
for _, part := range msg.Content {
if part.GetType() == ContentTypeToolResult {
count++
}
}
}
return count
}
func uniqueTransitionStepIndices(transitions []ReActTransition) []int {
seen := make(map[int]struct{})
order := make([]int, 0, len(transitions))
for _, transition := range transitions {
if _, ok := seen[transition.StepIndex]; ok {
continue
}
seen[transition.StepIndex] = struct{}{}
order = append(order, transition.StepIndex)
}
return order
}
func TestAgent_Generate_PropagatesCurrentStepIndex(t *testing.T) {
t.Parallel()
runtime := &recordingToolRuntime{}
observer := &captureObserver{}
tool := &mockTool{name: "echo", description: "echo", parameters: map[string]any{"type": "object"}}
agent := NewAgent(
&multiStepToolModel{},
WithTools(tool),
WithToolRuntime(runtime),
WithTransitionObserver(observer),
)
result, err := agent.Generate(t.Context(), AgentCall{Prompt: "run"})
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, []int{0, 1}, runtime.stepIndices)
assert.Equal(t, []int{0, 1, 2}, uniqueTransitionStepIndices(observer.Snapshot()))
}
func TestAgent_Stream_PropagatesCurrentStepIndex(t *testing.T) {
t.Parallel()
runtime := &recordingToolRuntime{}
observer := &captureObserver{}
tool := &mockTool{name: "echo", description: "echo", parameters: map[string]any{"type": "object"}}
agent := NewAgent(
&multiStepToolModel{},
WithTools(tool),
WithToolRuntime(runtime),
WithTransitionObserver(observer),
)
result, err := agent.Stream(t.Context(), AgentStreamCall{Prompt: "run"})
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, []int{0, 1}, runtime.stepIndices)
assert.Equal(t, []int{0, 1, 2}, uniqueTransitionStepIndices(observer.Snapshot()))
}

View file

@ -0,0 +1,23 @@
package fantasy
import "context"
type ctxStepIndexKey struct{}
// WithStepIndex annotates execution context with the current agent step.
func WithStepIndex(ctx context.Context, stepIndex int) context.Context {
return context.WithValue(ctx, ctxStepIndexKey{}, stepIndex)
}
// StepIndexFromCtx returns the current agent step carried in context.
// Missing values default to zero so callers can remain best-effort.
func StepIndexFromCtx(ctx context.Context) int {
v := ctx.Value(ctxStepIndexKey{})
if v == nil {
return 0
}
if i, ok := v.(int); ok {
return i
}
return 0
}

View file

@ -301,9 +301,9 @@ func NewRegistry(_ string) []app.Task {
NewCommandTask("eval", "Run the eval suite", evalRunSpecs, nil, nil), NewCommandTask("eval", "Run the eval suite", evalRunSpecs, nil, nil),
NewCommandTask("eval-fixtures", "Prepare eval fixture workspace", evalFixturesSpecs, nil, nil), NewCommandTask("eval-fixtures", "Prepare eval fixture workspace", evalFixturesSpecs, nil, nil),
NewCommandTask("eval-view", "Open the promptfoo results viewer", evalViewSpecs, nil, nil), NewCommandTask("eval-view", "Open the promptfoo results viewer", evalViewSpecs, nil, nil),
NewShellTask("eval-clean", "Cleanup eval artifacts", simpleScript("rm -rf eval/results eval/bin"), nil), NewShellTask("eval-clean", "Cleanup eval artifacts", evalCleanScript, nil),
NewShellTask("eval-compare", "Run A/B comparison of current branch vs main", simpleScript("cd eval && DEVCONTAINER_EXEC= EVAL_NPM_CMD=\"npx --yes\" ./scripts/compare.sh --repeat 3"), nil), NewShellTask("eval-compare", "Run A/B comparison of current branch vs main", evalCompareScript, nil),
NewShellTask("eval-test", "Run Go-native component evals", staticGoScript("-v ./eval/go_evals/..."), nil), NewShellTask("eval-test", "Run Go-native component evals", evalTestScript, nil),
} }
return tasks return tasks
} }
@ -711,14 +711,15 @@ func evalBuildSpecs(c *app.Context) []runner.CommandSpec {
version, commit, buildTime, goVersion, version, commit, buildTime, goVersion,
) )
return []runner.CommandSpec{ return []runner.CommandSpec{
{Name: goBinary, Args: []string{"generate", "./..."}}, {Name: goBinary, Args: []string{"generate", "./..."}, Dir: evalTaskRoot(c)},
{Name: "mkdir", Args: []string{"-p", "eval/bin"}}, {Name: "mkdir", Args: []string{"-p", "eval/bin"}, Dir: evalTaskRoot(c)},
{ {
Name: goBinary, Name: goBinary,
Args: append( Args: append(
append(append([]string{"build"}, goFlags...), "-ldflags", ldFlags), append(append([]string{"build"}, goFlags...), "-ldflags", ldFlags),
"-o", filepath.Join("eval", "bin", "eval-runner"), "./eval/cmd/eval-runner", "-o", filepath.Join("eval", "bin", "eval-runner"), "./eval/cmd/eval-runner",
), ),
Dir: evalTaskRoot(c),
}, },
} }
} }
@ -727,29 +728,31 @@ func evalRunSpecs(c *app.Context) []runner.CommandSpec {
cfgPath := cEnv(c, "DRAGONSCALE_EVAL_CONFIG", "./configs/default.json") cfgPath := cEnv(c, "DRAGONSCALE_EVAL_CONFIG", "./configs/default.json")
baseCfg := cEnv(c, "DRAGONSCALE_EVAL_BASE_CONFIG", "") baseCfg := cEnv(c, "DRAGONSCALE_EVAL_BASE_CONFIG", "")
debug := cEnv(c, "DRAGONSCALE_EVAL_DEBUG", "") != "" debug := cEnv(c, "DRAGONSCALE_EVAL_DEBUG", "") != ""
promptfooArgs := strings.Fields(cEnv(c, "DRAGONSCALE_PROMPTFOO_ARGS", "--no-cache --no-progress-bar")) promptfooArgs := strings.Fields(cEnv(c, "DRAGONSCALE_PROMPTFOO_ARGS", "--no-cache --no-progress-bar -j 1"))
if len(promptfooArgs) == 0 { if len(promptfooArgs) == 0 {
promptfooArgs = []string{"--no-cache", "--no-progress-bar"} promptfooArgs = []string{"--no-cache", "--no-progress-bar", "-j", "1"}
} }
var specs []runner.CommandSpec specs := append([]runner.CommandSpec{}, maybeEvalBuildSpecs(c)...)
if debug && strings.TrimSpace(baseCfg) != "" { if debug && strings.TrimSpace(baseCfg) != "" {
specs = append(specs, runner.CommandSpec{ specs = append(specs, runner.CommandSpec{
Name: "echo", Name: "echo",
Args: []string{fmt.Sprintf("DRAGONSCALE_EVAL_BASE_CONFIG=%s", baseCfg)}, Args: []string{fmt.Sprintf("DRAGONSCALE_EVAL_BASE_CONFIG=%s", baseCfg)},
Dir: evalTaskRoot(c),
}) })
} }
if debug { if debug {
specs = append(specs, runner.CommandSpec{ specs = append(specs, runner.CommandSpec{
Name: "echo", Name: "echo",
Args: []string{fmt.Sprintf("DRAGONSCALE_EVAL_CONFIG=%s", cfgPath)}, Args: []string{fmt.Sprintf("DRAGONSCALE_EVAL_CONFIG=%s", cfgPath)},
Dir: evalTaskRoot(c),
}) })
} }
args := append([]string{"--yes", "promptfoo", "eval", "--config", "promptfooconfig.yaml"}, promptfooArgs...) args := append([]string{"--yes", "promptfoo", "eval", "--config", "promptfooconfig.yaml"}, promptfooArgs...)
specs = append(specs, runner.CommandSpec{ specs = append(specs, runner.CommandSpec{
Name: "npx", Name: "npx",
Args: args, Args: args,
Dir: filepath.Join(c.Root, "eval"), Dir: evalWorkspaceDir(c),
Env: []string{fmt.Sprintf("DRAGONSCALE_EVAL_CONFIG=%s", cfgPath)}, Env: []string{fmt.Sprintf("DRAGONSCALE_EVAL_CONFIG=%s", cfgPath)},
}) })
return specs return specs
@ -774,16 +777,17 @@ func evalFixturesSpecs(c *app.Context) []runner.CommandSpec {
sourceFixture := filepath.Join("eval", "fixtures", "sample_data.txt") sourceFixture := filepath.Join("eval", "fixtures", "sample_data.txt")
specs := []runner.CommandSpec{ specs := []runner.CommandSpec{
{Name: "mkdir", Args: []string{"-p", sandbox}}, {Name: "mkdir", Args: []string{"-p", sandbox}, Dir: evalTaskRoot(c)},
{Name: "rm", Args: append([]string{"-f"}, files...)}, {Name: "rm", Args: append([]string{"-f"}, files...), Dir: evalTaskRoot(c)},
{Name: "rm", Args: []string{"-rf", project}}, {Name: "rm", Args: []string{"-rf", project}, Dir: evalTaskRoot(c)},
{Name: "mkdir", Args: []string{"-p", skills}}, {Name: "mkdir", Args: []string{"-p", skills}, Dir: evalTaskRoot(c)},
{Name: "bash", Args: []string{"-lc", fmt.Sprintf("printf '%%s\\n%%s\\n' \"dragonscale eval fixture — hello from the eval harness\" \"This is line two of the fixture file.\" > %q", fixture)}}, {Name: "bash", Args: []string{"-lc", fmt.Sprintf("printf '%%s\\n%%s\\n' \"dragonscale eval fixture — hello from the eval harness\" \"This is line two of the fixture file.\" > %q", fixture)}, Dir: evalTaskRoot(c)},
{Name: "cp", Args: []string{"-f", sourceFixture, shared}}, {Name: "cp", Args: []string{"-f", sourceFixture, shared}, Dir: evalTaskRoot(c)},
} }
specs = append(specs, runner.CommandSpec{ specs = append(specs, runner.CommandSpec{
Name: "bash", Name: "bash",
Args: []string{"-lc", "if [ -d eval/fixtures/skills ]; then cp -rf eval/fixtures/skills/. " + strconv.Quote(skills) + "; fi"}, Args: []string{"-lc", "if [ -d eval/fixtures/skills ]; then cp -rf eval/fixtures/skills/. " + strconv.Quote(skills) + "; fi"},
Dir: evalTaskRoot(c),
}) })
return specs return specs
} }
@ -793,7 +797,7 @@ func evalViewSpecs(c *app.Context) []runner.CommandSpec {
{ {
Name: "npx", Name: "npx",
Args: []string{"--yes", "promptfoo", "view"}, Args: []string{"--yes", "promptfoo", "view"},
Dir: filepath.Join(c.Root, "eval"), Dir: evalWorkspaceDir(c),
}, },
} }
} }
@ -805,3 +809,45 @@ func outputOrDefault(command string) string {
} }
return strings.TrimSpace(string(output)) return strings.TrimSpace(string(output))
} }
func evalTaskRoot(c *app.Context) string {
if c != nil && strings.TrimSpace(c.Root) != "" {
return c.Root
}
if wd, err := os.Getwd(); err == nil {
return wd
}
return "."
}
func evalWorkspaceDir(c *app.Context) string {
return filepath.Join(evalTaskRoot(c), "eval")
}
func evalRunnerBinaryPath(c *app.Context) string {
return filepath.Join(evalWorkspaceDir(c), "bin", "eval-runner")
}
func hasEvalSourceTree(c *app.Context) bool {
_, err := os.Stat(filepath.Join(evalWorkspaceDir(c), "cmd", "eval-runner", "main.go"))
return err == nil
}
func maybeEvalBuildSpecs(c *app.Context) []runner.CommandSpec {
if !hasEvalSourceTree(c) {
return nil
}
return evalBuildSpecs(c)
}
func evalCompareScript(c *app.Context) string {
return "cd " + shellSingleQuote(evalWorkspaceDir(c)) + " && DEVCONTAINER_EXEC= EVAL_NPM_CMD=\"npx --yes\" ./scripts/compare.sh --repeat 3\n"
}
func evalTestScript(c *app.Context) string {
return "cd " + shellSingleQuote(evalTaskRoot(c)) + " && $GO test -v ./eval/go_evals/...\n"
}
func evalCleanScript(c *app.Context) string {
return "rm -rf " + shellSingleQuote(filepath.Join(evalWorkspaceDir(c), "results")) + " " + shellSingleQuote(filepath.Join(evalWorkspaceDir(c), "bin")) + "\n"
}

View file

@ -207,7 +207,7 @@ func TestEvalRunSpecsPreservesEvalConfig(t *testing.T) {
require.Equal(t, "npx", specs[1].Name) require.Equal(t, "npx", specs[1].Name)
require.Contains(t, strings.Join(specs[1].Env, " "), "DRAGONSCALE_EVAL_CONFIG=./configs/override.json") require.Contains(t, strings.Join(specs[1].Env, " "), "DRAGONSCALE_EVAL_CONFIG=./configs/override.json")
require.NotContains(t, strings.Join(specs[1].Env, " "), "DRAGONSCALE_EVAL_CONFIG=./configs/default.json") require.NotContains(t, strings.Join(specs[1].Env, " "), "DRAGONSCALE_EVAL_CONFIG=./configs/default.json")
require.Contains(t, strings.Join(specs[1].Args, " "), "--yes promptfoo eval --config promptfooconfig.yaml --no-cache --no-progress-bar") require.Contains(t, strings.Join(specs[1].Args, " "), "--yes promptfoo eval --config promptfooconfig.yaml --no-cache --no-progress-bar -j 1")
compareScript, err := os.ReadFile(filepath.Clean(filepath.Join("..", "..", "..", "eval", "scripts", "compare.sh"))) compareScript, err := os.ReadFile(filepath.Clean(filepath.Join("..", "..", "..", "eval", "scripts", "compare.sh")))
require.NoError(t, err) require.NoError(t, err)
@ -216,6 +216,8 @@ func TestEvalRunSpecsPreservesEvalConfig(t *testing.T) {
require.Contains(t, content, "trap cleanup_compare_config EXIT INT TERM") require.Contains(t, content, "trap cleanup_compare_config EXIT INT TERM")
require.Contains(t, content, "DRAGONSCALE_EVAL_CONFIG: \"${EVAL_CONFIG}\"") require.Contains(t, content, "DRAGONSCALE_EVAL_CONFIG: \"${EVAL_CONFIG}\"")
require.NotContains(t, content, "DRAGONSCALE_EVAL_HOST_HOME: \"/host_home\"") require.NotContains(t, content, "DRAGONSCALE_EVAL_HOST_HOME: \"/host_home\"")
require.Contains(t, content, "PROMPTFOO_ARGS=\"${DRAGONSCALE_PROMPTFOO_ARGS:---no-cache --no-progress-bar -j 1}\"")
require.Contains(t, content, "git -C \"$PROJECT_ROOT\" worktree add --force --detach")
require.Contains(t, content, "TEMP_CONFIG") require.Contains(t, content, "TEMP_CONFIG")
} }
@ -235,7 +237,7 @@ func TestEvalRunSpecsUsesBaseConfigWhenSetAndDebugEnabled(t *testing.T) {
require.Equal(t, "echo", specs[1].Name) require.Equal(t, "echo", specs[1].Name)
require.Equal(t, []string{"DRAGONSCALE_EVAL_CONFIG=./configs/default.json"}, specs[1].Args) require.Equal(t, []string{"DRAGONSCALE_EVAL_CONFIG=./configs/default.json"}, specs[1].Args)
require.Equal(t, "npx", specs[2].Name) require.Equal(t, "npx", specs[2].Name)
require.Equal(t, []string{"--yes", "promptfoo", "eval", "--config", "promptfooconfig.yaml", "--no-cache", "--no-progress-bar"}, specs[2].Args) require.Equal(t, []string{"--yes", "promptfoo", "eval", "--config", "promptfooconfig.yaml", "--no-cache", "--no-progress-bar", "-j", "1"}, specs[2].Args)
} }
func TestEvalRunSpecsUsesPromptfooArgsOverride(t *testing.T) { func TestEvalRunSpecsUsesPromptfooArgsOverride(t *testing.T) {
@ -284,7 +286,8 @@ func TestEvalCompareTaskDisablesNestedDevcontainerExecution(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
require.Len(t, fake.Calls, 1) require.Len(t, fake.Calls, 1)
script := strings.Join(fake.Calls[0].Args, " ") script := strings.Join(fake.Calls[0].Args, " ")
require.Contains(t, script, "cd eval && DEVCONTAINER_EXEC= EVAL_NPM_CMD=\"npx --yes\" ./scripts/compare.sh --repeat 3") require.Contains(t, script, evalWorkspaceDir(ctx))
require.Contains(t, script, "DEVCONTAINER_EXEC= EVAL_NPM_CMD=\"npx --yes\" ./scripts/compare.sh --repeat 3")
joinedEnv := strings.Join(fake.Calls[0].Env, " ") joinedEnv := strings.Join(fake.Calls[0].Env, " ")
require.Contains(t, joinedEnv, "DEVCONTAINER_EXEC=npx --yes @devcontainers/cli exec --workspace-folder \"$PWD\" --") require.Contains(t, joinedEnv, "DEVCONTAINER_EXEC=npx --yes @devcontainers/cli exec --workspace-folder \"$PWD\" --")
@ -293,6 +296,8 @@ func TestEvalCompareTaskDisablesNestedDevcontainerExecution(t *testing.T) {
content := string(compareScript) content := string(compareScript)
require.Contains(t, content, "export DEVCONTAINER_EXEC=\"\"") require.Contains(t, content, "export DEVCONTAINER_EXEC=\"\"")
require.Contains(t, content, "make DEVCONTAINER_EXEC= eval-build") require.Contains(t, content, "make DEVCONTAINER_EXEC= eval-build")
require.Contains(t, content, "git -C \"$PROJECT_ROOT\" worktree add --force --detach")
require.Contains(t, content, "git -C \"$PROJECT_ROOT\" worktree remove --force")
} }
func TestEvalCompareTaskPassesBaseConfigEnvToRunner(t *testing.T) { func TestEvalCompareTaskPassesBaseConfigEnvToRunner(t *testing.T) {
@ -450,6 +455,67 @@ func TestEvalViewTaskRunsInEvalDirectory(t *testing.T) {
require.Equal(t, []string{"--yes", "promptfoo", "view"}, fake.Calls[0].Args) require.Equal(t, []string{"--yes", "promptfoo", "view"}, fake.Calls[0].Args)
} }
func TestEvalRunSpecsPrependsBuildWhenRunnerMissingAndSourceTreeExists(t *testing.T) {
t.Parallel()
root := t.TempDir()
require.NoError(t, os.MkdirAll(filepath.Join(root, "eval", "cmd", "eval-runner"), 0o755))
require.NoError(t, os.WriteFile(filepath.Join(root, "eval", "cmd", "eval-runner", "main.go"), []byte("package main\nfunc main() {}\n"), 0o644))
specs := evalRunSpecs(&app.Context{Root: root})
require.GreaterOrEqual(t, len(specs), 4)
require.Equal(t, "go", specs[0].Name)
require.Equal(t, []string{"generate", "./..."}, specs[0].Args)
require.Equal(t, root, specs[0].Dir)
require.Equal(t, "npx", specs[len(specs)-1].Name)
require.Equal(t, filepath.Join(root, "eval"), specs[len(specs)-1].Dir)
}
func TestEvalRunSpecsPrependsBuildWhenRunnerAlreadyExists(t *testing.T) {
t.Parallel()
root := t.TempDir()
require.NoError(t, os.MkdirAll(filepath.Join(root, "eval", "cmd", "eval-runner"), 0o755))
require.NoError(t, os.WriteFile(filepath.Join(root, "eval", "cmd", "eval-runner", "main.go"), []byte("package main\nfunc main() {}\n"), 0o644))
require.NoError(t, os.MkdirAll(filepath.Join(root, "eval", "bin"), 0o755))
require.NoError(t, os.WriteFile(filepath.Join(root, "eval", "bin", "eval-runner"), []byte("stale-binary"), 0o755))
specs := evalRunSpecs(&app.Context{Root: root})
require.GreaterOrEqual(t, len(specs), 4)
require.Equal(t, "go", specs[0].Name)
require.Equal(t, []string{"generate", "./..."}, specs[0].Args)
require.Equal(t, root, specs[0].Dir)
require.Equal(t, "npx", specs[len(specs)-1].Name)
}
func TestEvalTasksUseRepoRootPathsWhenCwdIsNested(t *testing.T) {
t.Parallel()
root := t.TempDir()
nested := filepath.Join(root, "pkg", "agent")
require.NoError(t, os.MkdirAll(nested, 0o755))
ctx := &app.Context{
Root: root,
Cwd: nested,
}
fake := &runner.FakeRunner{Result: runner.CommandResult{ExitCode: 0}}
evalView := findTaskByName(t, NewRegistry(ctx.Root), "eval-view")
_, err := evalView.Run(context.Background(), fake, ctx)
require.NoError(t, err)
require.Len(t, fake.Calls, 1)
require.Equal(t, filepath.Join(root, "eval"), fake.Calls[0].Dir)
fake.Calls = nil
evalCompare := findTaskByName(t, NewRegistry(ctx.Root), "eval-compare")
_, err = evalCompare.Run(context.Background(), fake, ctx)
require.NoError(t, err)
require.Len(t, fake.Calls, 1)
compareCmd := strings.Join(fake.Calls[0].Args, " ")
require.Contains(t, compareCmd, filepath.Join(root, "eval"))
require.Contains(t, compareCmd, "DEVCONTAINER_EXEC= EVAL_NPM_CMD=\"npx --yes\" ./scripts/compare.sh --repeat 3")
}
func TestEvalFixturesTaskCreatesFixtureCommands(t *testing.T) { func TestEvalFixturesTaskCreatesFixtureCommands(t *testing.T) {
ctx := &app.Context{ ctx := &app.Context{
Root: t.TempDir(), Root: t.TempDir(),

View file

@ -0,0 +1,560 @@
package agent
import (
"context"
"fmt"
"strings"
"time"
jsonv2 "github.com/go-json-experiment/json"
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/observation"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
memstore "github.com/ZanzyTHEbar/dragonscale/pkg/memory/store"
"github.com/ZanzyTHEbar/dragonscale/pkg/messages"
)
const projectionImmutableHistoryLimit = 2048
type projectionBudget struct {
Total int
System int
Recent int
DAG int
Retrieval int
}
func newProjectionBudget(maxTokens int) projectionBudget {
if maxTokens <= 0 {
maxTokens = 4096
}
system := maxTokens * 35 / 100
recent := maxTokens * 40 / 100
dagBudget := maxTokens * 10 / 100
retrieval := maxTokens - system - recent - dagBudget
if retrieval < 0 {
retrieval = 0
}
return projectionBudget{
Total: maxTokens,
System: system,
Recent: recent,
DAG: dagBudget,
Retrieval: retrieval,
}
}
type TurnContextBuildRequest struct {
ProjectionRequest memory.ProjectionRequest
CurrentMessage string
NoHistory bool
FallbackHistory []messages.Message
Summary string
}
type TurnContextBuild struct {
Projection *memory.ActiveContextProjection
History []messages.Message
}
// DefaultActiveContextBuilder materializes the contract-defined projection
// from the runtime's existing system prompt builder, immutable history, DAG
// summaries, and hybrid memory retrieval.
type DefaultActiveContextBuilder struct {
agentID string
contextBuilder *ContextBuilder
sessions sessionSummaryReader
memDelegate memory.MemoryDelegate
memoryStore *memstore.MemoryStore
queries *memsqlc.Queries
}
type sessionSummaryReader interface {
GetSummary(key string) string
}
var _ memory.ActiveContextBuilder = (*DefaultActiveContextBuilder)(nil)
func NewDefaultActiveContextBuilder(agentID string, contextBuilder *ContextBuilder, sessions sessionSummaryReader, memDelegate memory.MemoryDelegate, memoryStore *memstore.MemoryStore, queries *memsqlc.Queries) *DefaultActiveContextBuilder {
return &DefaultActiveContextBuilder{
agentID: agentID,
contextBuilder: contextBuilder,
sessions: sessions,
memDelegate: memDelegate,
memoryStore: memoryStore,
queries: queries,
}
}
func (b *DefaultActiveContextBuilder) BuildActiveContext(ctx context.Context, req memory.ProjectionRequest) (*memory.ActiveContextProjection, error) {
built, err := b.BuildTurnContext(ctx, TurnContextBuildRequest{
ProjectionRequest: req,
})
if err != nil {
return nil, err
}
return built.Projection, nil
}
func (b *DefaultActiveContextBuilder) BuildTurnContext(ctx context.Context, req TurnContextBuildRequest) (*TurnContextBuild, error) {
agentID := strings.TrimSpace(req.ProjectionRequest.AgentID)
if agentID == "" {
agentID = b.agentID
}
budget := newProjectionBudget(req.ProjectionRequest.MaxTokens)
now := time.Now().UTC()
projection := &memory.ActiveContextProjection{
AgentID: agentID,
SessionKey: req.ProjectionRequest.SessionKey,
BudgetTokens: budget.Total,
GeneratedAt: now,
ProjectionRef: fmt.Sprintf("%s:%d", req.ProjectionRequest.SessionKey, now.UnixNano()),
}
summary := strings.TrimSpace(req.Summary)
if summary == "" && b.sessions != nil {
summary = strings.TrimSpace(b.sessions.GetSummary(req.ProjectionRequest.SessionKey))
}
systemSegments := b.buildSystemSegments(req.ProjectionRequest.SessionKey, summary, budget.System)
projection.Segments = append(projection.Segments, systemSegments...)
var immutableHistory []*memory.ImmutableMessage
if !req.NoHistory && b.memDelegate != nil {
historyRows, err := b.memDelegate.ListImmutableMessages(ctx, req.ProjectionRequest.SessionKey, projectionImmutableHistoryLimit, 0)
if err != nil {
logger.WarnCF("context", "Active context builder failed to load immutable history",
map[string]interface{}{
"session_key": req.ProjectionRequest.SessionKey,
"error": err.Error(),
})
} else {
immutableHistory = historyRows
}
}
history, historySegments := b.buildHistorySegments(req.ProjectionRequest.SessionKey, immutableHistory, req.FallbackHistory, budget.Recent)
projection.Segments = append(projection.Segments, historySegments...)
if !req.NoHistory {
projection.Segments = append(projection.Segments, b.buildDAGSegments(ctx, req.ProjectionRequest.SessionKey, immutableHistory, budget.DAG)...)
projection.Segments = append(projection.Segments, b.buildRetrievalSegments(ctx, req.ProjectionRequest.SessionKey, req.CurrentMessage, budget.Retrieval)...)
}
return &TurnContextBuild{
Projection: projection,
History: history,
}, nil
}
func (b *DefaultActiveContextBuilder) buildSystemSegments(sessionKey, summary string, budget int) []memory.ProjectionSegment {
if b.contextBuilder == nil || budget <= 0 {
return nil
}
candidates := make([]memory.ProjectionSegment, 0, 2)
systemPrompt := strings.TrimSpace(b.contextBuilder.BuildSystemPromptWithBudget(0))
if systemPrompt != "" {
candidates = append(candidates, memory.ProjectionSegment{
Kind: memory.ProjectionSegmentSystem,
Source: "runtime_system",
Text: systemPrompt,
Tokens: observation.EstimateTokens(systemPrompt),
Ref: zeroSpanRef(sessionKey),
})
}
if summary != "" {
summaryText := "## Summary of Previous Conversation\n\n" + summary
candidates = append(candidates, memory.ProjectionSegment{
Kind: memory.ProjectionSegmentSystem,
Source: "session_summary",
Text: summaryText,
Tokens: observation.EstimateTokens(summaryText),
Ref: zeroSpanRef(sessionKey),
})
}
return fitSegmentsToBudget(candidates, budget)
}
func (b *DefaultActiveContextBuilder) buildHistorySegments(sessionKey string, immutableHistory []*memory.ImmutableMessage, fallbackHistory []messages.Message, budget int) ([]messages.Message, []memory.ProjectionSegment) {
if budget <= 0 {
return nil, nil
}
if len(immutableHistory) > 0 {
return buildHistoryFromImmutable(sessionKey, immutableHistory, budget)
}
return buildHistoryFromFallback(sessionKey, fallbackHistory, budget)
}
func buildHistoryFromImmutable(sessionKey string, immutableHistory []*memory.ImmutableMessage, budget int) ([]messages.Message, []memory.ProjectionSegment) {
if len(immutableHistory) == 0 {
return nil, nil
}
minTail := dag.TailMessageCount(budget)
selected := make([]int, 0, minTail)
remaining := budget
for idx := len(immutableHistory) - 1; idx >= 0; idx-- {
msg := immutableHistory[idx]
tokens := msg.TokenEstimate
if tokens <= 0 {
tokens = observation.EstimateTokens(renderImmutableMessageText(msg))
}
if len(selected) >= minTail && remaining-tokens < 0 {
break
}
selected = append(selected, idx)
remaining -= tokens
}
reverseInts(selected)
history := make([]messages.Message, 0, len(selected))
segments := make([]memory.ProjectionSegment, 0, len(selected))
for _, idx := range selected {
msg := immutableHistory[idx]
sessionMsg := immutableToSessionMessage(msg)
history = append(history, sessionMsg)
text := renderSessionMessageText(sessionMsg)
kind := memory.ProjectionSegmentRecent
if msg.Role == "tool" {
kind = memory.ProjectionSegmentTool
}
segments = append(segments, memory.ProjectionSegment{
Kind: kind,
Source: "immutable:" + msg.Role,
Text: text,
Tokens: max(msg.TokenEstimate, observation.EstimateTokens(text)),
Ref: immutableMessageRef(sessionKey, idx, msg),
})
}
return history, segments
}
func buildHistoryFromFallback(sessionKey string, fallbackHistory []messages.Message, budget int) ([]messages.Message, []memory.ProjectionSegment) {
if len(fallbackHistory) == 0 {
return nil, nil
}
minTail := dag.TailMessageCount(budget)
selected := make([]int, 0, minTail)
remaining := budget
for idx := len(fallbackHistory) - 1; idx >= 0; idx-- {
msg := fallbackHistory[idx]
tokens := observation.EstimateTokens(renderSessionMessageText(msg))
if len(selected) >= minTail && remaining-tokens < 0 {
break
}
selected = append(selected, idx)
remaining -= tokens
}
reverseInts(selected)
history := make([]messages.Message, 0, len(selected))
segments := make([]memory.ProjectionSegment, 0, len(selected))
for _, idx := range selected {
msg := fallbackHistory[idx]
history = append(history, msg)
kind := memory.ProjectionSegmentRecent
if msg.Role == "tool" {
kind = memory.ProjectionSegmentTool
}
segments = append(segments, memory.ProjectionSegment{
Kind: kind,
Source: "session_cache:" + msg.Role,
Text: renderSessionMessageText(msg),
Tokens: observation.EstimateTokens(renderSessionMessageText(msg)),
Ref: memory.ImmutableSpanRef{
SessionKey: sessionKey,
StartIdx: idx,
EndIdx: idx + 1,
},
})
}
return history, segments
}
func (b *DefaultActiveContextBuilder) buildDAGSegments(ctx context.Context, sessionKey string, immutableHistory []*memory.ImmutableMessage, budget int) []memory.ProjectionSegment {
if budget <= 0 || b.queries == nil {
return nil
}
snapshot, err := b.queries.GetLatestDAGSnapshotBySession(ctx, memsqlc.GetLatestDAGSnapshotBySessionParams{
AgentID: b.agentID,
SessionKey: sessionKey,
})
if err != nil {
return nil
}
nodes, err := b.queries.ListDAGNodesBySnapshotID(ctx, memsqlc.ListDAGNodesBySnapshotIDParams{
SnapshotID: snapshot.ID,
})
if err != nil || len(nodes) == 0 {
return nil
}
selected := selectDAGNodesForBudget(nodes, budget)
segments := make([]memory.ProjectionSegment, 0, len(selected))
for _, node := range selected {
text := fmt.Sprintf("[%d-%d] %s", node.StartIdx, max(node.EndIdx-1, node.StartIdx), strings.TrimSpace(node.Summary))
segments = append(segments, memory.ProjectionSegment{
Kind: memory.ProjectionSegmentDAG,
Source: fmt.Sprintf("dag:%s:%s", sessionKey, node.NodeID),
Text: text,
Tokens: max(int(node.Tokens), observation.EstimateTokens(text)),
Ref: dagNodeRef(sessionKey, immutableHistory, node),
})
}
return fitSegmentsToBudget(segments, budget)
}
func (b *DefaultActiveContextBuilder) buildRetrievalSegments(ctx context.Context, sessionKey, currentMessage string, budget int) []memory.ProjectionSegment {
if budget <= 0 || b.memoryStore == nil || strings.TrimSpace(currentMessage) == "" {
return nil
}
results, err := b.memoryStore.Search(ctx, currentMessage, memory.SearchOptions{
AgentID: b.agentID,
SessionKey: sessionKey,
Limit: 8,
})
if err != nil {
logger.WarnCF("context", "Active context builder retrieval search failed",
map[string]interface{}{
"session_key": sessionKey,
"error": err.Error(),
})
return nil
}
segments := make([]memory.ProjectionSegment, 0, len(results))
seen := make(map[string]struct{}, len(results))
for _, result := range results {
if strings.TrimSpace(result.Content) == "" {
continue
}
if strings.HasPrefix(result.Source, "working-context:") || strings.HasPrefix(result.Source, "dag:") {
continue
}
kind := memory.ProjectionSegmentArchival
if result.Source == sessionKey {
kind = memory.ProjectionSegmentRecall
}
key := string(kind) + ":" + result.ID.String() + ":" + result.Source
if _, ok := seen[key]; ok {
continue
}
seen[key] = struct{}{}
segments = append(segments, memory.ProjectionSegment{
Kind: kind,
Source: result.Source,
Text: strings.TrimSpace(result.Content),
Tokens: observation.EstimateTokens(result.Content),
Ref: zeroSpanRef(sessionKey),
})
}
return fitSegmentsToBudget(segments, budget)
}
func selectDAGNodesForBudget(nodes []memsqlc.DagNode, budget int) []memsqlc.DagNode {
if budget <= 0 || len(nodes) == 0 {
return nil
}
grouped := map[int64][]memsqlc.DagNode{}
totals := map[int64]int{}
for _, node := range nodes {
grouped[node.Level] = append(grouped[node.Level], node)
totals[node.Level] += max(int(node.Tokens), observation.EstimateTokens(node.Summary))
}
for _, level := range []int64{int64(dag.LevelSession), int64(dag.LevelSection), int64(dag.LevelChunk)} {
if len(grouped[level]) == 0 {
continue
}
if totals[level] <= budget {
return grouped[level]
}
}
for _, level := range []int64{int64(dag.LevelSession), int64(dag.LevelSection), int64(dag.LevelChunk)} {
if len(grouped[level]) > 0 {
return grouped[level]
}
}
return nil
}
func fitSegmentsToBudget(segments []memory.ProjectionSegment, budget int) []memory.ProjectionSegment {
if budget <= 0 || len(segments) == 0 {
return nil
}
out := make([]memory.ProjectionSegment, 0, len(segments))
remaining := budget
for _, seg := range segments {
if strings.TrimSpace(seg.Text) == "" {
continue
}
if seg.Tokens <= 0 {
seg.Tokens = observation.EstimateTokens(seg.Text)
}
if seg.Tokens <= remaining {
out = append(out, seg)
remaining -= seg.Tokens
continue
}
if remaining < 32 {
break
}
seg.Text = truncateToTokenBudget(seg.Text, remaining)
seg.Tokens = observation.EstimateTokens(seg.Text)
out = append(out, seg)
break
}
return out
}
func immutableToSessionMessage(msg *memory.ImmutableMessage) messages.Message {
if msg == nil {
return messages.Message{}
}
sessionMsg := messages.Message{
Role: msg.Role,
Content: msg.Content,
ToolCallID: msg.ToolCallID,
}
if strings.TrimSpace(msg.ToolCalls) != "" {
_ = jsonv2.Unmarshal([]byte(msg.ToolCalls), &sessionMsg.ToolCalls)
}
return sessionMsg
}
func renderImmutableMessageText(msg *memory.ImmutableMessage) string {
if msg == nil {
return ""
}
return renderSessionMessageText(immutableToSessionMessage(msg))
}
func renderSessionMessageText(msg messages.Message) string {
content := strings.TrimSpace(msg.Content)
if content != "" {
return content
}
if len(msg.ToolCalls) == 0 {
return ""
}
parts := make([]string, 0, len(msg.ToolCalls))
for _, call := range msg.ToolCalls {
name := strings.TrimSpace(call.Name)
if call.Function != nil && strings.TrimSpace(call.Function.Name) != "" {
name = strings.TrimSpace(call.Function.Name)
}
if name == "" {
name = "tool"
}
parts = append(parts, name)
}
return "Tool calls: " + strings.Join(parts, ", ")
}
func zeroSpanRef(sessionKey string) memory.ImmutableSpanRef {
return memory.ImmutableSpanRef{SessionKey: sessionKey}
}
func immutableMessageRef(sessionKey string, idx int, msg *memory.ImmutableMessage) memory.ImmutableSpanRef {
ref := memory.ImmutableSpanRef{
SessionKey: sessionKey,
StartIdx: idx,
EndIdx: idx + 1,
}
if msg == nil {
return ref
}
ref.FirstID = msg.ID
ref.LastID = msg.ID
ref.FromTime = msg.CreatedAt
ref.ToTime = msg.CreatedAt
return ref
}
func dagNodeRef(sessionKey string, immutableHistory []*memory.ImmutableMessage, node memsqlc.DagNode) memory.ImmutableSpanRef {
ref := memory.ImmutableSpanRef{
SessionKey: sessionKey,
StartIdx: max(0, int(node.StartIdx)),
EndIdx: max(0, int(node.EndIdx)),
}
if len(immutableHistory) == 0 {
if ref.EndIdx < ref.StartIdx {
ref.EndIdx = ref.StartIdx
}
return ref
}
if ref.StartIdx > len(immutableHistory) {
ref.StartIdx = len(immutableHistory)
}
if ref.EndIdx > len(immutableHistory) {
ref.EndIdx = len(immutableHistory)
}
if ref.EndIdx < ref.StartIdx {
ref.EndIdx = ref.StartIdx
}
if ref.EndIdx == ref.StartIdx {
return ref
}
first := immutableHistory[ref.StartIdx]
last := immutableHistory[ref.EndIdx-1]
ref.FirstID = first.ID
ref.LastID = last.ID
ref.FromTime = first.CreatedAt
ref.ToTime = last.CreatedAt
return ref
}
func reverseInts(values []int) {
for i, j := 0, len(values)-1; i < j; i, j = i+1, j-1 {
values[i], values[j] = values[j], values[i]
}
}

View file

@ -0,0 +1,168 @@
package agent
import (
"os"
"path/filepath"
"testing"
pkgroot "github.com/ZanzyTHEbar/dragonscale/pkg"
"github.com/ZanzyTHEbar/dragonscale/pkg/bus"
"github.com/ZanzyTHEbar/dragonscale/pkg/config"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAssembleContext_UsesActiveContextProjection(t *testing.T) {
t.Parallel()
tmpDir, err := os.MkdirTemp("", "active-context-projection-*")
require.NoError(t, err)
defer os.RemoveAll(tmpDir)
cfg := config.DefaultConfig()
cfg.Agents.Defaults.Sandbox = tmpDir
cfg.Agents.Defaults.MaxTokens = 4096
cfg.Agents.Defaults.MaxToolIterations = 4
cfg.Memory.DBPath = filepath.Join(tmpDir, "active-context.db")
al := mustNewAgentLoop(t, cfg, bus.NewMessageBus(), newMockLanguageModel("ok"))
require.NotNil(t, al.activeContextBuilder)
sessionKey := "projection-session"
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "user",
Content: "How does the runtime work?",
TokenEstimate: 16,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "assistant",
Content: "It runs a ReAct loop with persisted state.",
TokenEstimate: 18,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "tool",
Content: "{\"result\":\"stored\"}",
ToolCallID: "tool-1",
TokenEstimate: 12,
})
ac, err := al.assembleContext(t.Context(), processOptions{
SessionKey: sessionKey,
UserMessage: "Summarize the prior tool result",
EnableSummary: false,
})
require.NoError(t, err)
require.NotNil(t, ac.projection)
assert.True(t, ac.projection.HasLosslessRefs())
assert.NotEmpty(t, ac.fantasyHistory)
assert.Len(t, ac.fantasyHistory, 3)
assert.Contains(t, projectionKinds(ac.projection), memory.ProjectionSegmentSystem)
assert.Contains(t, projectionKinds(ac.projection), memory.ProjectionSegmentRecent)
assert.Contains(t, projectionKinds(ac.projection), memory.ProjectionSegmentTool)
assert.Contains(t, ac.systemPrompt, "You have access to the following tools")
}
func TestActiveContextBuilder_IncludesPersistedDAGProjection(t *testing.T) {
t.Parallel()
tmpDir, err := os.MkdirTemp("", "active-context-dag-*")
require.NoError(t, err)
defer os.RemoveAll(tmpDir)
cfg := config.DefaultConfig()
cfg.Agents.Defaults.Sandbox = tmpDir
cfg.Agents.Defaults.MaxTokens = 4096
cfg.Memory.DBPath = filepath.Join(tmpDir, "active-context-dag.db")
al := mustNewAgentLoop(t, cfg, bus.NewMessageBus(), newMockLanguageModel("ok"))
require.NotNil(t, al.activeContextBuilder)
sessionKey := "dag-projection-session"
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "user",
Content: "one",
TokenEstimate: 4,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "assistant",
Content: "two",
TokenEstimate: 4,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "user",
Content: "three",
TokenEstimate: 4,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "assistant",
Content: "four",
TokenEstimate: 4,
})
persister, ok := al.memDelegate.(dag.DAGPersister)
require.True(t, ok)
tree := dag.NewDAG()
tree.Add(&dag.Node{
ID: "session-root",
Level: dag.LevelSession,
Summary: "compressed session summary",
Tokens: 24,
StartIdx: 0,
EndIdx: 4,
})
tree.SetRoots([]string{"session-root"})
require.NoError(t, persister.PersistDAG(t.Context(), pkgroot.NAME, sessionKey, &dag.PersistSnapshot{
FromMsgIdx: 0,
ToMsgIdx: 4,
MsgCount: 4,
DAG: tree,
}))
built, err := al.activeContextBuilder.BuildTurnContext(t.Context(), TurnContextBuildRequest{
ProjectionRequest: memory.ProjectionRequest{
AgentID: pkgroot.NAME,
SessionKey: sessionKey,
MaxTokens: 4096,
},
CurrentMessage: "summarize",
})
require.NoError(t, err)
require.NotNil(t, built.Projection)
assert.Contains(t, projectionKinds(built.Projection), memory.ProjectionSegmentDAG)
}
func insertImmutableMessage(t *testing.T, al *AgentLoop, msg *memory.ImmutableMessage) {
t.Helper()
require.NoError(t, al.memDelegate.InsertImmutableMessage(t.Context(), msg))
}
func projectionKinds(projection *memory.ActiveContextProjection) []memory.ProjectionSegmentKind {
if projection == nil {
return nil
}
kinds := make([]memory.ProjectionSegmentKind, 0, len(projection.Segments))
for _, seg := range projection.Segments {
kinds = append(kinds, seg.Kind)
}
return kinds
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,284 @@
package agent
import (
"context"
"fmt"
"strings"
"time"
fantasy "charm.land/fantasy"
"github.com/ZanzyTHEbar/dragonscale/pkg"
"github.com/ZanzyTHEbar/dragonscale/pkg/agent/conversations"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
"github.com/ZanzyTHEbar/dragonscale/pkg/messages"
jsonv2 "github.com/go-json-experiment/json"
)
const checkpointTurnCompletedEvent = "turn.completed"
type sessionMessageLister interface {
ListSessionMessages(ctx context.Context, agentID, sessionKey, role string, limit int) ([]*memory.RecallItem, error)
}
func checkpointNameForRun(runID ids.UUID) string {
return fmt.Sprintf("run-%s-complete", runID.String())
}
func (al *AgentLoop) persistRunCheckpoint(ctx context.Context, opts processOptions, metrics agentRunMetrics) {
if al == nil || al.stateStore == nil || al.queries == nil {
return
}
if opts.ConversationID.IsZero() || opts.RunID.IsZero() {
return
}
history := al.sessions.GetHistory(opts.SessionKey)
snapshot := conversations.NewCheckpointSnapshot(
opts.SessionKey,
opts.ConversationID.String(),
opts.RunID.String(),
checkpointTurnCompletedEvent,
metrics.StepCount,
history,
)
meta := map[string]any{
"event": checkpointTurnCompletedEvent,
"session_key": opts.SessionKey,
"conversation_id": opts.ConversationID.String(),
"run_id": opts.RunID.String(),
"step_count": metrics.StepCount,
"tool_calls": metrics.ToolCalls,
"errors": metrics.Errors,
}
runState, err := al.stateStore.AddRunState(ctx, opts.RunID, metrics.StepCount, fantasy.ReActStateDone, snapshot)
if err != nil {
logger.WarnCF("agent", "Failed to persist checkpointable run snapshot", map[string]any{
"session_key": opts.SessionKey,
"run_id": opts.RunID.String(),
"error": err.Error(),
})
} else {
checkpointName := checkpointNameForRun(opts.RunID)
meta["checkpoint_name"] = checkpointName
meta["run_state_id"] = runState.ID.String()
checkpointStore := NewCheckpointStore(al.queries)
if _, err := checkpointStore.CreateCheckpoint(ctx, opts.ConversationID, checkpointName, runState.ID, meta); err != nil {
logger.WarnCF("agent", "Failed to create runtime checkpoint", map[string]any{
"session_key": opts.SessionKey,
"run_id": opts.RunID.String(),
"checkpoint": checkpointName,
"error": err.Error(),
})
}
}
if _, err := al.stateStore.UpdateRunStatus(ctx, opts.RunID, "completed", meta); err != nil {
logger.WarnCF("agent", "Failed to update run completion status", map[string]any{
"session_key": opts.SessionKey,
"run_id": opts.RunID.String(),
"error": err.Error(),
})
}
}
func (al *AgentLoop) RestoreSessionFromCheckpoint(ctx context.Context, sessionKey, checkpointName string) error {
sessionKey = strings.TrimSpace(sessionKey)
if sessionKey == "" {
return fmt.Errorf("session key is required")
}
checkpointName = strings.TrimSpace(checkpointName)
if checkpointName == "" {
return fmt.Errorf("checkpoint name is required")
}
conversationID, err := al.lookupConversationIDForSession(ctx, sessionKey)
if err != nil {
return err
}
store := conversations.New(al.queries)
_, snapshot, err := store.LoadCheckpointSnapshot(ctx, conversationID, checkpointName)
if err != nil {
return err
}
if err := al.replaceSessionHistory(ctx, sessionKey, snapshot.Messages); err != nil {
return err
}
al.conversationIDs.Store(sessionKey, conversationID)
return nil
}
func (al *AgentLoop) ForkSessionFromCheckpoint(ctx context.Context, sourceSessionKey, checkpointName, forkSessionKey string) (ids.UUID, error) {
sourceSessionKey = strings.TrimSpace(sourceSessionKey)
if sourceSessionKey == "" {
return ids.UUID{}, fmt.Errorf("source session key is required")
}
checkpointName = strings.TrimSpace(checkpointName)
if checkpointName == "" {
return ids.UUID{}, fmt.Errorf("checkpoint name is required")
}
forkSessionKey = strings.TrimSpace(forkSessionKey)
if forkSessionKey == "" {
return ids.UUID{}, fmt.Errorf("fork session key is required")
}
conversationID, err := al.lookupConversationIDForSession(ctx, sourceSessionKey)
if err != nil {
return ids.UUID{}, err
}
store := conversations.New(al.queries)
title := forkSessionKey
conv, err := store.ForkFromCheckpoint(ctx, conversations.ForkFromCheckpointParams{
FromConversationID: conversationID.String(),
CheckpointName: checkpointName,
Title: &title,
})
if err != nil {
return ids.UUID{}, err
}
_, snapshot, err := store.LoadCheckpointSnapshot(ctx, conversationID, checkpointName)
if err != nil {
return ids.UUID{}, err
}
if err := al.replaceSessionHistory(ctx, forkSessionKey, snapshot.Messages); err != nil {
return ids.UUID{}, err
}
al.conversationIDs.Store(forkSessionKey, conv.ID)
return conv.ID, nil
}
func (al *AgentLoop) lookupConversationIDForSession(ctx context.Context, sessionKey string) (ids.UUID, error) {
if al == nil || al.queries == nil {
return ids.UUID{}, fmt.Errorf("runtime persistence is not initialized")
}
sessionKey = strings.TrimSpace(sessionKey)
if sessionKey == "" {
return ids.UUID{}, fmt.Errorf("session key is required")
}
if cached, ok := al.conversationIDs.Load(sessionKey); ok {
return cached, nil
}
conv, err := al.queries.GetLatestAgentConversationByTitle(ctx, memsqlc.GetLatestAgentConversationByTitleParams{Title: &sessionKey})
if err != nil {
return ids.UUID{}, fmt.Errorf("lookup conversation for session %q: %w", sessionKey, err)
}
al.conversationIDs.Store(sessionKey, conv.ID)
return conv.ID, nil
}
func (al *AgentLoop) replaceSessionHistory(ctx context.Context, sessionKey string, history []messages.Message) error {
if al == nil || al.sessions == nil {
return fmt.Errorf("session manager is not initialized")
}
restored := conversations.HydrationMessages(history, conversations.MaxCheckpointHydrationMessages)
if lister, ok := al.memDelegate.(sessionMessageLister); ok && al.memDelegate != nil {
limit := 32
if count, err := al.memDelegate.CountRecallItems(ctx, pkg.NAME, sessionKey); err == nil && count+32 > limit {
limit = count + 32
}
if rows, err := lister.ListSessionMessages(ctx, pkg.NAME, sessionKey, "", limit); err == nil {
for _, item := range rows {
if err := al.memDelegate.SoftDeleteRecallItem(ctx, pkg.NAME, item.ID); err != nil {
logger.WarnCF("agent", "Failed to suppress prior session message during checkpoint restore", map[string]any{
"session_key": sessionKey,
"message_id": item.ID.String(),
"error": err.Error(),
})
}
}
} else {
logger.WarnCF("agent", "Failed to list prior session messages during checkpoint restore", map[string]any{
"session_key": sessionKey,
"error": err.Error(),
})
}
}
al.sessions.ReplaceHistory(sessionKey, restored, "")
if session := al.sessions.GetOrCreate(sessionKey); session != nil {
session.Messages = conversations.HydrationMessages(restored, len(restored))
session.Summary = ""
session.Updated = time.Now().UTC()
}
for _, msg := range restored {
if err := al.persistCheckpointSessionMessage(ctx, sessionKey, msg); err != nil {
return err
}
}
al.contextTreeCache.Delete(sessionKey)
return al.sessions.Save(sessionKey)
}
func (al *AgentLoop) persistCheckpointSessionMessage(ctx context.Context, sessionKey string, msg messages.Message) error {
if al.memDelegate == nil {
return nil
}
now := time.Now().UTC()
item := &memory.RecallItem{
ID: ids.New(),
AgentID: pkg.NAME,
SessionKey: sessionKey,
Role: msg.Role,
Sector: memory.SectorEpisodic,
Importance: 0.5,
Salience: 0.5,
DecayRate: 0.01,
Content: msg.Content,
Tags: "session-message",
CreatedAt: now,
UpdatedAt: now,
}
if err := al.memDelegate.InsertRecallItem(ctx, item); err != nil {
return fmt.Errorf("persist restored recall item: %w", err)
}
immutableWriter, ok := al.memDelegate.(interface {
InsertImmutableMessage(ctx context.Context, msg *memory.ImmutableMessage) error
})
if !ok {
return nil
}
imMsg := &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: msg.Role,
Content: msg.Content,
ToolCallID: msg.ToolCallID,
ToolCalls: checkpointToolCallsJSON(msg),
TokenEstimate: estimateCheckpointMessageTokens(msg.Content),
}
if err := immutableWriter.InsertImmutableMessage(ctx, imMsg); err != nil {
return fmt.Errorf("persist restored immutable message: %w", err)
}
return nil
}
func estimateCheckpointMessageTokens(content string) int {
return (len(content) + 3) / 4
}
func checkpointToolCallsJSON(msg messages.Message) string {
if len(msg.ToolCalls) == 0 {
return ""
}
b, err := jsonv2.Marshal(msg.ToolCalls)
if err != nil {
return ""
}
return string(b)
}

View file

@ -0,0 +1,188 @@
package agent
import (
"context"
"os"
"path/filepath"
"testing"
"time"
pkgroot "github.com/ZanzyTHEbar/dragonscale/pkg"
"github.com/ZanzyTHEbar/dragonscale/pkg/agent/conversations"
"github.com/ZanzyTHEbar/dragonscale/pkg/bus"
"github.com/ZanzyTHEbar/dragonscale/pkg/config"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
"github.com/ZanzyTHEbar/dragonscale/pkg/messages"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestIntegration_RuntimeCheckpoint_PersistsSnapshotAndCheckpoint(t *testing.T) {
t.Parallel()
al := newCheckpointTestAgentLoop(t, "checkpoint reply")
ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
defer cancel()
sessionKey := "checkpoint-session"
response, err := al.processMessage(ctx, bus.InboundMessage{
Channel: "test",
SenderID: "user1",
ChatID: "chat1",
Content: "hello checkpoint runtime",
SessionKey: sessionKey,
})
require.NoError(t, err)
assert.Contains(t, response, "checkpoint reply")
require.NoError(t, al.sessions.Save(sessionKey))
run, checkpoint, snapshot := loadCheckpointFixture(t, al, sessionKey)
assert.Equal(t, "completed", run.Status)
assert.Equal(t, checkpointTurnCompletedEvent, snapshot.Event)
assert.NotEmpty(t, snapshot.Messages)
assert.Equal(t, "user", snapshot.Messages[0].Role)
assert.Equal(t, "hello checkpoint runtime", snapshot.Messages[0].Content)
assert.Equal(t, "assistant", snapshot.Messages[len(snapshot.Messages)-1].Role)
assert.Contains(t, snapshot.Messages[len(snapshot.Messages)-1].Content, "checkpoint reply")
assert.Equal(t, checkpointNameForRun(run.ID), checkpoint.Name)
}
func TestAgentLoop_RestoreSessionFromCheckpoint_ReplacesPersistedHistory(t *testing.T) {
t.Parallel()
al := newCheckpointTestAgentLoop(t, "restored checkpoint reply")
ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
defer cancel()
sessionKey := "restore-session"
_, err := al.processMessage(ctx, bus.InboundMessage{
Channel: "test",
SenderID: "user1",
ChatID: "chat1",
Content: "restore me",
SessionKey: sessionKey,
})
require.NoError(t, err)
require.NoError(t, al.sessions.Save(sessionKey))
run, checkpoint, snapshot := loadCheckpointFixture(t, al, sessionKey)
require.False(t, run.ID.IsZero())
al.sessions.AddMessage(sessionKey, "user", "mutated user")
al.sessions.AddMessage(sessionKey, "assistant", "mutated assistant")
require.NoError(t, al.sessions.Save(sessionKey))
err = al.RestoreSessionFromCheckpoint(ctx, sessionKey, checkpoint.Name)
require.NoError(t, err)
expected := conversations.HydrationMessages(snapshot.Messages, conversations.MaxCheckpointHydrationMessages)
assert.Equal(t, checkpointHistoryView(expected), checkpointHistoryView(al.sessions.GetHistory(sessionKey)))
lister, ok := al.memDelegate.(sessionMessageLister)
require.True(t, ok)
rows, err := lister.ListSessionMessages(ctx, pkgroot.NAME, sessionKey, "", 64)
require.NoError(t, err)
assert.Equal(t, checkpointHistoryView(expected), recallHistoryView(rows))
}
func TestAgentLoop_ForkSessionFromCheckpoint_CreatesHydratedChildSession(t *testing.T) {
t.Parallel()
al := newCheckpointTestAgentLoop(t, "fork checkpoint reply")
ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
defer cancel()
sessionKey := "fork-source-session"
_, err := al.processMessage(ctx, bus.InboundMessage{
Channel: "test",
SenderID: "user1",
ChatID: "chat1",
Content: "fork me",
SessionKey: sessionKey,
})
require.NoError(t, err)
require.NoError(t, al.sessions.Save(sessionKey))
_, checkpoint, snapshot := loadCheckpointFixture(t, al, sessionKey)
forkSessionKey := "fork-child-session"
childConversationID, err := al.ForkSessionFromCheckpoint(ctx, sessionKey, checkpoint.Name, forkSessionKey)
require.NoError(t, err)
require.False(t, childConversationID.IsZero())
expected := conversations.HydrationMessages(snapshot.Messages, conversations.MaxCheckpointHydrationMessages)
assert.Equal(t, checkpointHistoryView(expected), checkpointHistoryView(al.sessions.GetHistory(forkSessionKey)))
seeded, err := al.queries.ListAgentMessagesByConversationID(ctx, memsqlc.ListAgentMessagesByConversationIDParams{
ConversationID: childConversationID,
})
require.NoError(t, err)
assert.Equal(t, checkpointHistoryView(expected), agentMessageHistoryView(seeded))
}
func newCheckpointTestAgentLoop(t *testing.T, response string) *AgentLoop {
t.Helper()
tmpDir, err := os.MkdirTemp("", "agent-checkpoint-*")
require.NoError(t, err)
t.Cleanup(func() { _ = os.RemoveAll(tmpDir) })
cfg := config.DefaultConfig()
cfg.Agents.Defaults.Sandbox = tmpDir
cfg.Agents.Defaults.MaxTokens = 4096
cfg.Agents.Defaults.MaxToolIterations = 4
cfg.Memory.DBPath = filepath.Join(tmpDir, "agent-checkpoint.db")
return mustNewAgentLoop(t, cfg, bus.NewMessageBus(), newMockLanguageModel(response))
}
func loadCheckpointFixture(t *testing.T, al *AgentLoop, sessionKey string) (memsqlc.AgentRun, memsqlc.AgentCheckpoint, conversations.CheckpointSnapshot) {
t.Helper()
conversationID, err := al.lookupConversationIDForSession(t.Context(), sessionKey)
require.NoError(t, err)
run, err := al.queries.GetLatestAgentRunByConversationID(t.Context(), memsqlc.GetLatestAgentRunByConversationIDParams{
ConversationID: conversationID,
})
require.NoError(t, err)
checkpoint := mustGetCheckpoint(t, al, conversationID, checkpointNameForRun(run.ID))
runState, err := al.queries.GetAgentRunStateByID(t.Context(), memsqlc.GetAgentRunStateByIDParams{ID: checkpoint.RunStateID})
require.NoError(t, err)
snapshot, err := conversations.DecodeCheckpointSnapshot(runState.SnapshotJson)
require.NoError(t, err)
return run, checkpoint, snapshot
}
func mustGetCheckpoint(t *testing.T, al *AgentLoop, conversationID ids.UUID, checkpointName string) memsqlc.AgentCheckpoint {
t.Helper()
checkpoint, err := NewCheckpointStore(al.queries).GetCheckpoint(t.Context(), conversationID, checkpointName)
require.NoError(t, err)
return checkpoint
}
func checkpointHistoryView(history []messages.Message) []string {
out := make([]string, 0, len(history))
for _, msg := range history {
out = append(out, msg.Role+":"+msg.Content)
}
return out
}
func recallHistoryView(rows []*memory.RecallItem) []string {
out := make([]string, 0, len(rows))
for _, row := range rows {
out = append(out, row.Role+":"+row.Content)
}
return out
}
func agentMessageHistoryView(rows []memsqlc.AgentMessage) []string {
out := make([]string, 0, len(rows))
for _, row := range rows {
out = append(out, row.Role+":"+row.Content)
}
return out
}

View file

@ -31,6 +31,7 @@ type ContextBuilder struct {
knowledgeBlock string // Pre-rendered knowledge block from Focus completions knowledgeBlock string // Pre-rendered knowledge block from Focus completions
contextTreeBlock string // Pre-rendered Context-Tree selected history contextTreeBlock string // Pre-rendered Context-Tree selected history
contextWindow int // Max tokens for context window (0 = no limit) contextWindow int // Max tokens for context window (0 = no limit)
sessionKeyFn func() string // Active session resolver for session-scoped prompt sections
cacheMu sync.Mutex cacheMu sync.Mutex
skillsCache string skillsCache string
@ -108,6 +109,10 @@ func (cb *ContextBuilder) SetContextWindow(tokens int) {
cb.contextWindow = tokens cb.contextWindow = tokens
} }
func (cb *ContextBuilder) SetSessionResolver(sessionKeyFn func() string) {
cb.sessionKeyFn = sessionKeyFn
}
func (cb *ContextBuilder) getIdentity() string { func (cb *ContextBuilder) getIdentity() string {
now := time.Now().Format("2006-01-02 15:04 (Monday)") now := time.Now().Format("2006-01-02 15:04 (Monday)")
workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace)) workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace))
@ -128,7 +133,7 @@ You are dragonscale, a helpful AI assistant.
## Workspace ## Workspace
Your workspace is at: %s Your workspace is at: %s
- Skills: %s/skills/{skill-name}/SKILL.md - Skills are loaded through skill_search / skill_read. Do NOT infer SKILL.md paths or read skills via read_file.
%s %s
@ -144,8 +149,25 @@ Your workspace is at: %s
5. **Completion discipline** - For actionable requests, execute the required tools before your final answer. Do NOT end with only intent statements like "I'll do that" or "let me do that." 5. **Completion discipline** - For actionable requests, execute the required tools before your final answer. Do NOT end with only intent statements like "I'll do that" or "let me do that."
6. **Context Management** - You MUST consolidate your context to stay effective during long tasks. Use start_focus at the beginning of any investigation or multi-step task. After 10-15 tool calls, call complete_focus with a summary of what you learned and accomplished. This compresses your working context and persists knowledge for future reference. Failing to consolidate will degrade your performance as context grows.`, 6. **Plans vs actions** - If the user only wants a plan, schedule, explanation, summary, or workflow, answer directly without tools unless they explicitly ask you to persist, modify, search, or execute something.
now, runtime, workspacePath, workspacePath, toolsSection)
7. **Direct tool routing** - When the task clearly maps to a tool, call that tool directly instead of using tool_search or tool_call first.
- Skills: use skill_search to discover skills, then skill_read to load one by name. Do NOT use tool_search for skill discovery.
- Files: use read_file, write_file, edit_file, append_file, and list_dir directly. Do NOT use shell redirection for normal file edits.
- If the task says replace, edit, patch, or update existing text, prefer edit_file over write_file.
- If the task says append or add to the end of a file, prefer append_file over write_file or exec.
- Shell: use exec with the raw command only, e.g. {"command":"uname -s"}. Keep working_dir separate; never mix paths or commentary into command.
- Commitments: use memory to capture/store commitments, deadlines, decisions, and follow-ups. Use obligation only when the user wants an actual tracked reminder lifecycle.
8. **Exact argument discipline** - Use the tool's exact parameter names. Examples:
- edit_file => {"path":"edit_target.txt","old_text":"world","new_text":"dragonscale"}
- append_file => {"path":"append_test.txt","content":"line two\n"}
- skill_read => {"name":"eval-test-skill"}
9. **Stop when verified** - After the requested change is completed and a verification read/result confirms success, stop calling tools and answer the user. Do NOT repeat the same write/edit/read cycle.
10. **Context Management** - You MUST consolidate your context to stay effective during long tasks. Use start_focus at the beginning of any investigation or multi-step task. After 10-15 tool calls, call complete_focus with a summary of what you learned and accomplished. This compresses your working context and persists knowledge for future reference. Failing to consolidate will degrade your performance as context grows.`,
now, runtime, workspacePath, toolsSection)
} }
func (cb *ContextBuilder) buildToolsSection() string { func (cb *ContextBuilder) buildToolsSection() string {
@ -177,6 +199,10 @@ type contextSection struct {
} }
func (cb *ContextBuilder) BuildSystemPrompt() string { func (cb *ContextBuilder) BuildSystemPrompt() string {
return cb.BuildSystemPromptWithBudget(cb.tokenBudgetTokens())
}
func (cb *ContextBuilder) BuildSystemPromptWithBudget(budgetTokens int) string {
// Collect sections in priority order // Collect sections in priority order
sections := []contextSection{} sections := []contextSection{}
@ -234,7 +260,6 @@ Do NOT assume skill content — always load before applying.
// token budget proportional to its priority weight. Surplus from small // token budget proportional to its priority weight. Surplus from small
// sections redistributes to higher-priority ones. Sections that still // sections redistributes to higher-priority ones. Sections that still
// exceed their allocation are truncated rather than dropped entirely. // exceed their allocation are truncated rather than dropped entirely.
budgetTokens := cb.tokenBudgetTokens()
totalTokens := 0 totalTokens := 0
sectionTokens := make([]int, len(sections)) sectionTokens := make([]int, len(sections))
for i, s := range sections { for i, s := range sections {
@ -267,6 +292,57 @@ Do NOT assume skill content — always load before applying.
return prompt return prompt
} }
func (cb *ContextBuilder) RenderProjection(projection *memory.ActiveContextProjection, channel, chatID string) string {
if projection == nil {
return cb.BuildSystemPrompt()
}
sections := make([]string, 0, len(projection.Segments)+1)
for _, seg := range projection.Segments {
if strings.TrimSpace(seg.Text) == "" {
continue
}
switch seg.Kind {
case memory.ProjectionSegmentSystem:
sections = append(sections, seg.Text)
case memory.ProjectionSegmentDAG:
sections = append(sections, "## Compressed Session Context\n\n"+seg.Text)
case memory.ProjectionSegmentRecall:
sections = append(sections, "## Recall Memory\n\n"+seg.Text)
case memory.ProjectionSegmentArchival:
sections = append(sections, "## Archival Memory\n\n"+seg.Text)
}
}
if channel != "" && chatID != "" {
sections = append(sections, fmt.Sprintf("## Current Session\nChannel: %s\nChat ID: %s", channel, chatID))
}
systemPrompt := strings.Join(sections, "\n\n---\n\n")
if systemPrompt == "" {
systemPrompt = cb.BuildSystemPrompt()
}
logger.DebugCF("agent", "System prompt built",
map[string]interface{}{
"total_chars": len(systemPrompt),
"total_lines": strings.Count(systemPrompt, "\n") + 1,
"section_count": strings.Count(systemPrompt, "\n\n---\n\n") + 1,
})
preview := systemPrompt
if len(preview) > 500 {
preview = preview[:500] + "... (truncated)"
}
logger.DebugCF("agent", "System prompt preview",
map[string]interface{}{
"preview": preview,
})
return systemPrompt
}
// tokenBudgetTokens returns the maximum token count for the system prompt, // tokenBudgetTokens returns the maximum token count for the system prompt,
// derived from the context window size. Returns 0 if no limit is configured. // derived from the context window size. Returns 0 if no limit is configured.
func (cb *ContextBuilder) tokenBudgetTokens() int { func (cb *ContextBuilder) tokenBudgetTokens() int {
@ -444,8 +520,15 @@ func (cb *ContextBuilder) buildWorkingContextSection() string {
var parts []string var parts []string
sessionKey := "default"
if cb.sessionKeyFn != nil {
if resolved := strings.TrimSpace(cb.sessionKeyFn()); resolved != "" {
sessionKey = resolved
}
}
// Inject working context (hot tier) // Inject working context (hot tier)
wc, err := cb.memoryStore.GetWorkingContext(ctx, pkg.NAME, "default") wc, err := cb.memoryStore.GetWorkingContext(ctx, pkg.NAME, sessionKey)
if err == nil && wc != "" { if err == nil && wc != "" {
parts = append(parts, "## Working Context\n\n"+wc) parts = append(parts, "## Working Context\n\n"+wc)
} }

View file

@ -0,0 +1,26 @@
package agent
import (
"strings"
"testing"
)
func TestSystemPromptIncludesDirectToolRoutingHints(t *testing.T) {
cb := NewContextBuilder(t.TempDir())
prompt := cb.BuildSystemPromptWithBudget(0)
expectedSnippets := []string{
"Plans vs actions",
"Direct tool routing",
"Do NOT infer SKILL.md paths",
"use skill_search to discover skills",
"use memory to capture/store commitments",
"{\"command\":\"uname -s\"}",
}
for _, snippet := range expectedSnippets {
if !strings.Contains(prompt, snippet) {
t.Fatalf("expected prompt to contain %q", snippet)
}
}
}

View file

@ -0,0 +1,106 @@
package agent
import (
"context"
"fmt"
"path/filepath"
"sync"
"testing"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
memstore "github.com/ZanzyTHEbar/dragonscale/pkg/memory/store"
"github.com/ZanzyTHEbar/dragonscale/pkg/messages"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestApplyContextTreeSelection_UsesSemanticEmbeddings(t *testing.T) {
t.Parallel()
target := "SQL latency issue on write path"
query := "database stalls"
al := &AgentLoop{
contextBuilder: NewContextBuilder(filepath.Join(t.TempDir(), "workspace")),
contextWindow: 4096,
memoryStore: memstore.New(nil, nil, fixtureEmbedder{query: {1, 0}, target: {1, 0}}, memstore.Config{ContextWindowTokens: 4096}),
contextTreeCache: sync.Map{},
}
history := semanticSelectionHistory(target)
tail := al.applyContextTreeSelection(t.Context(), "semantic-session", query, history)
require.NotEmpty(t, tail)
assert.Contains(t, al.contextBuilder.contextTreeBlock, target)
}
func TestApplyContextTreeSelection_PreservesSelectionAccessCounts(t *testing.T) {
t.Parallel()
target := "SQL latency issue on write path"
query := "database stalls"
al := &AgentLoop{
contextBuilder: NewContextBuilder(filepath.Join(t.TempDir(), "workspace")),
contextWindow: 4096,
memoryStore: memstore.New(nil, nil, fixtureEmbedder{query: {1, 0}, target: {1, 0}}, memstore.Config{ContextWindowTokens: 4096}),
contextTreeCache: sync.Map{},
}
history := semanticSelectionHistory(target)
_ = al.applyContextTreeSelection(t.Context(), "semantic-session", query, history)
firstCached, ok := al.contextTreeCache.Load("semantic-session")
require.True(t, ok)
firstEntry := firstCached.(contextTreeCacheEntry)
require.NotEmpty(t, firstEntry.selectedKeys)
firstKey := firstEntry.selectedKeys[0]
firstCount := firstEntry.accessCounts[firstKey]
_ = al.applyContextTreeSelection(t.Context(), "semantic-session", query, history)
secondCached, ok := al.contextTreeCache.Load("semantic-session")
require.True(t, ok)
secondEntry := secondCached.(contextTreeCacheEntry)
require.NotEmpty(t, secondEntry.selectedKeys)
assert.Greater(t, secondEntry.accessCounts[firstKey], firstCount)
}
type fixtureEmbedder map[string]memory.Embedding
func (f fixtureEmbedder) Embed(_ context.Context, text string) (memory.Embedding, error) {
if embedding, ok := f[text]; ok {
return embedding, nil
}
return memory.Embedding{0, 1}, nil
}
func (f fixtureEmbedder) EmbedBatch(ctx context.Context, texts []string) ([]memory.Embedding, error) {
out := make([]memory.Embedding, 0, len(texts))
for _, text := range texts {
embedding, err := f.Embed(ctx, text)
if err != nil {
return nil, err
}
out = append(out, embedding)
}
return out, nil
}
func (f fixtureEmbedder) Dimensions() int { return 2 }
func (f fixtureEmbedder) Model() string { return "fixture" }
func semanticSelectionHistory(target string) []messages.Message {
history := make([]messages.Message, 0, 40)
for i := 0; i < 40; i++ {
content := fmt.Sprintf("routine message %02d about unrelated topic", i)
if i == 5 {
content = target
}
role := "user"
if i%2 == 1 {
role = "assistant"
}
history = append(history, messages.Message{
Role: role,
Content: content,
})
}
return history
}

View file

@ -0,0 +1,70 @@
package conversations
import (
jsonv2 "github.com/go-json-experiment/json"
"github.com/ZanzyTHEbar/dragonscale/pkg/messages"
)
const MaxCheckpointHydrationMessages = 200
// CheckpointSnapshot is the canonical runtime payload persisted into
// agent_run_states for checkpoint hydration.
type CheckpointSnapshot struct {
SessionKey string `json:"session_key,omitempty"`
ConversationID string `json:"conversation_id,omitempty"`
RunID string `json:"run_id,omitempty"`
Event string `json:"event,omitempty"`
StepCount int `json:"step_count,omitempty"`
Messages []messages.Message `json:"messages"`
}
func NewCheckpointSnapshot(sessionKey, conversationID, runID, event string, stepCount int, history []messages.Message) CheckpointSnapshot {
return CheckpointSnapshot{
SessionKey: sessionKey,
ConversationID: conversationID,
RunID: runID,
Event: event,
StepCount: stepCount,
Messages: cloneCheckpointMessages(history),
}
}
func DecodeCheckpointSnapshot(raw []byte) (CheckpointSnapshot, error) {
if len(raw) == 0 {
return CheckpointSnapshot{}, nil
}
var snap CheckpointSnapshot
if err := jsonv2.Unmarshal(raw, &snap); err != nil {
return CheckpointSnapshot{}, err
}
snap.Messages = cloneCheckpointMessages(snap.Messages)
return snap, nil
}
func HydrationMessages(history []messages.Message, limit int) []messages.Message {
if limit <= 0 {
limit = MaxCheckpointHydrationMessages
}
if len(history) > limit {
history = history[len(history)-limit:]
}
return cloneCheckpointMessages(history)
}
func cloneCheckpointMessages(history []messages.Message) []messages.Message {
if len(history) == 0 {
return nil
}
cloned := make([]messages.Message, len(history))
copy(cloned, history)
for i := range cloned {
if len(cloned[i].ToolCalls) == 0 {
continue
}
toolCalls := make([]messages.ToolCall, len(cloned[i].ToolCalls))
copy(toolCalls, cloned[i].ToolCalls)
cloned[i].ToolCalls = toolCalls
}
return cloned
}

View file

@ -153,34 +153,11 @@ func (s *Store) ForkFromCheckpoint(ctx context.Context, p ForkFromCheckpointPara
return sqlc.AgentConversation{}, dserrors.New(dserrors.CodeInvalidArgument, "checkpoint_name is required") return sqlc.AgentConversation{}, dserrors.New(dserrors.CodeInvalidArgument, "checkpoint_name is required")
} }
cp, err := s.q.GetAgentCheckpointByConversationIDAndName(ctx, cp, snap, err := s.LoadCheckpointSnapshot(ctx, fromID, cpName)
sqlc.GetAgentCheckpointByConversationIDAndNameParams{
ConversationID: fromID,
Name: cpName,
})
if err != nil { if err != nil {
return sqlc.AgentConversation{}, err return sqlc.AgentConversation{}, err
} }
runState, err := s.q.GetAgentRunStateByID(ctx, sqlc.GetAgentRunStateByIDParams{ID: cp.RunStateID})
if err != nil {
return sqlc.AgentConversation{}, err
}
type msgSnapshot struct {
Role string `json:"role"`
Content string `json:"content"`
}
type snapshot struct {
Messages []msgSnapshot `json:"messages"`
}
var snap snapshot
if len(runState.SnapshotJson) > 0 {
if err := jsonv2.Unmarshal(runState.SnapshotJson, &snap); err != nil {
return sqlc.AgentConversation{}, dserrors.Wrapf(dserrors.CodeInternal, err, "parse snapshot for run state %s", cp.RunStateID)
}
}
conv, err := s.q.CreateAgentConversation(ctx, sqlc.CreateAgentConversationParams{ conv, err := s.q.CreateAgentConversation(ctx, sqlc.CreateAgentConversationParams{
ID: ids.New(), ID: ids.New(),
Title: p.Title, Title: p.Title,
@ -204,10 +181,7 @@ func (s *Store) ForkFromCheckpoint(ctx context.Context, p ForkFromCheckpointPara
}) })
// Seed messages from snapshot — cap to 200 to prevent pathological snapshots. // Seed messages from snapshot — cap to 200 to prevent pathological snapshots.
msgs := snap.Messages msgs := HydrationMessages(snap.Messages, MaxCheckpointHydrationMessages)
if len(msgs) > 200 {
msgs = msgs[len(msgs)-200:]
}
seedMeta := map[string]any{ seedMeta := map[string]any{
"seeded_from_conversation_id": fromID.String(), "seeded_from_conversation_id": fromID.String(),
@ -217,7 +191,7 @@ func (s *Store) ForkFromCheckpoint(ctx context.Context, p ForkFromCheckpointPara
seedMetaJSON, _ := jsonv2.Marshal(seedMeta) seedMetaJSON, _ := jsonv2.Marshal(seedMeta)
for _, m := range msgs { for _, m := range msgs {
if m.Role != "user" && m.Role != "assistant" { if !isCheckpointHydrationRole(m.Role) {
continue continue
} }
if strings.TrimSpace(m.Content) == "" { if strings.TrimSpace(m.Content) == "" {
@ -235,6 +209,44 @@ func (s *Store) ForkFromCheckpoint(ctx context.Context, p ForkFromCheckpointPara
return conv, nil return conv, nil
} }
func (s *Store) LoadCheckpointSnapshot(ctx context.Context, conversationID ids.UUID, checkpointName string) (sqlc.AgentCheckpoint, CheckpointSnapshot, error) {
if conversationID.IsZero() {
return sqlc.AgentCheckpoint{}, CheckpointSnapshot{}, dserrors.New(dserrors.CodeInvalidArgument, "conversation_id is required")
}
if strings.TrimSpace(checkpointName) == "" {
return sqlc.AgentCheckpoint{}, CheckpointSnapshot{}, dserrors.New(dserrors.CodeInvalidArgument, "checkpoint_name is required")
}
cp, err := s.q.GetAgentCheckpointByConversationIDAndName(ctx, sqlc.GetAgentCheckpointByConversationIDAndNameParams{
ConversationID: conversationID,
Name: strings.TrimSpace(checkpointName),
})
if err != nil {
return sqlc.AgentCheckpoint{}, CheckpointSnapshot{}, err
}
runState, err := s.q.GetAgentRunStateByID(ctx, sqlc.GetAgentRunStateByIDParams{ID: cp.RunStateID})
if err != nil {
return sqlc.AgentCheckpoint{}, CheckpointSnapshot{}, err
}
snap, err := DecodeCheckpointSnapshot(runState.SnapshotJson)
if err != nil {
return sqlc.AgentCheckpoint{}, CheckpointSnapshot{}, dserrors.Wrapf(dserrors.CodeInternal, err, "parse snapshot for run state %s", cp.RunStateID)
}
return cp, snap, nil
}
func isCheckpointHydrationRole(role string) bool {
switch role {
case "user", "assistant", "tool", "system":
return true
default:
return false
}
}
// ─── Merge ──────────────────────────────────────────────────────────────────── // ─── Merge ────────────────────────────────────────────────────────────────────
// MergeAsLinkedContextParams configures the MergeAsLinkedContext operation. // MergeAsLinkedContextParams configures the MergeAsLinkedContext operation.

View file

@ -0,0 +1,296 @@
package agent
import (
"errors"
"os"
"path/filepath"
"strings"
"testing"
fantasy "charm.land/fantasy"
"github.com/ZanzyTHEbar/dragonscale/pkg/skills"
"github.com/ZanzyTHEbar/dragonscale/pkg/tools"
)
func TestGroundFinalContentConditionalBranch(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
steps := []fantasy.StepResult{
stepWithToolResults(
toolText("read_file", "sample_data.txt contains the word fox"),
toolText("write_file", "found fox"),
),
}
got := al.groundFinalContent(
"Read the file sample_data.txt. If it contains the word 'fox', write 'found fox' to result.txt. Otherwise write 'no fox'.",
"Done.",
steps,
)
want := `The file contains "fox", so I wrote "found fox" to result.txt.`
if got != want {
t.Fatalf("unexpected grounded content\nwant: %q\ngot: %q", want, got)
}
}
func TestGroundFinalContentAddsReplacementReadback(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
steps := []fantasy.StepResult{
stepWithToolResults(
toolText("read_file", "hello dragonscale"),
),
}
got := al.groundFinalContent(
"First write a file called edit_target.txt with 'hello world'. Then edit it to replace 'world' with 'dragonscale'. Read it back and confirm.",
"Updated the file.",
steps,
)
if !strings.Contains(got, "dragonscale") {
t.Fatalf("expected grounded content to mention dragonscale, got %q", got)
}
if !strings.Contains(got, "Updated content: hello dragonscale") {
t.Fatalf("expected grounded content to include readback, got %q", got)
}
}
func TestGroundFinalContentFallsBackToReplacementConfirmation(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
steps := []fantasy.StepResult{
stepWithToolResults(
toolText("read_file", ","),
),
}
got := al.groundFinalContent(
"First write a file called edit_target.txt with 'hello world'. Then edit it to replace 'world' with 'dragonscale'. Read it back and confirm.",
"Updated the file.",
steps,
)
if !strings.Contains(got, "Confirmed replacement includes dragonscale.") {
t.Fatalf("expected generic replacement confirmation, got %q", got)
}
}
func TestGroundFinalContentAddsOSConfirmation(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
steps := []fantasy.StepResult{
stepWithToolResults(
toolText("exec", "Linux\n"),
toolText("read_file", "Linux\n"),
),
}
got := al.groundFinalContent(
"Run 'uname -s' to get the OS name, then write the result to a file called os_name.txt, then read it back and confirm.",
"Done! Here's what happened.",
steps,
)
if !strings.Contains(got, "Confirmed OS name: Linux.") {
t.Fatalf("expected OS grounding, got %q", got)
}
}
func TestGroundFinalContentAddsGenericReadBackConfirmation(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
steps := []fantasy.StepResult{
stepWithToolResults(
toolText("read_file", "step one step two"),
),
}
got := al.groundFinalContent(
"Create a file called chain_test.txt with 'step one'. Then append ' step two' to it. Finally read it back and tell me the full contents.",
"Done.",
steps,
)
if !strings.Contains(got, "Read-back confirmation: step one step two") {
t.Fatalf("expected read-back grounding, got %q", got)
}
}
func TestGroundFinalContentExpandsCommitmentPlan(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
got := al.groundFinalContent(
"I have three commitments: submit tax documents by March 15, follow up with Alex in 2 days, and renew my passport next month. Capture these commitments and give me a reminder/follow-up plan with explicit timing.",
"Captured the commitments.",
nil,
)
for _, snippet := range []string{
"submit tax documents by March 15",
"follow up with Alex in 2 days",
"renew my passport next month",
"Reminder/follow-up plan",
} {
if !strings.Contains(strings.ToLower(got), strings.ToLower(snippet)) {
t.Fatalf("expected grounded commitments to include %q, got %q", snippet, got)
}
}
}
func TestGroundFinalContentExpandsExactCommitmentRegister(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
got := al.groundFinalContent(
"Track these commitments exactly: send rent receipt tonight, book vet appointment tomorrow, and submit sprint notes by Friday. Return a commitment register and verification checklist.",
"Captured the commitments.",
nil,
)
for _, snippet := range []string{
"rent receipt",
"vet appointment",
"sprint notes",
} {
if !strings.Contains(strings.ToLower(got), strings.ToLower(snippet)) {
t.Fatalf("expected grounded commitments to include %q, got %q", snippet, got)
}
}
}
func TestGroundFinalContentRecoversSkillSummary(t *testing.T) {
t.Parallel()
skillsDir := t.TempDir()
skillDir := filepath.Join(skillsDir, "eval-test-skill")
if err := os.MkdirAll(skillDir, 0o755); err != nil {
t.Fatalf("mkdir skill dir: %v", err)
}
content := `---
name: eval-test-skill
description: greeting templates
tags: [eval, greeting]
domain: testing
---
# Eval Test Skill
- **Formal**: "Good day, {name}. How may I assist you?"
- **Casual**: "Hey {name}! What's up?"
- **Technical**: "Hello {name}, ready to debug some code?"
`
if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644); err != nil {
t.Fatalf("write skill file: %v", err)
}
loader := skills.NewSkillsLoader(skillsDir, "", "")
registry := tools.NewToolRegistry()
registry.Register(tools.NewSkillReadTool(loader))
al := &AgentLoop{tools: registry}
got := al.groundFinalContent(
"Read the 'eval-test-skill' skill and tell me what greeting templates it provides.",
"Let me try a cleaner approach:",
[]fantasy.StepResult{stepWithToolResults(toolText("skill_read", "name is required"))},
)
for _, snippet := range []string{"greeting templates", "Good day", "Hey", "debug some code"} {
if !strings.Contains(got, snippet) {
t.Fatalf("expected recovered skill summary to include %q, got %q", snippet, got)
}
}
}
func TestObservedToolNameUnwrapsToolCall(t *testing.T) {
t.Parallel()
got := observedToolName("tool_call", `{"tool_name":"read_file","arguments":{"path":"x"}}`)
if got != "read_file" {
t.Fatalf("expected nested tool name, got %q", got)
}
}
func TestResolveFinalContentPrefersToolResultOverPreamble(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
got, err := al.resolveFinalContent("", []fantasy.StepResult{
stepWithTextAndToolResults(
"Now I'll search the memory for that term:",
toolText("memory", "No results found for: xyzzy_nonexistent_topic_42"),
),
})
if err != nil {
t.Fatalf("resolveFinalContent returned error: %v", err)
}
if got != "No results found for: xyzzy_nonexistent_topic_42" {
t.Fatalf("expected tool result recovery, got %q", got)
}
}
func TestGroundFinalContentOverridesContradictoryExecSuccess(t *testing.T) {
t.Parallel()
al := &AgentLoop{}
got := al.groundFinalContent(
"Run the command 'sleep 120' and tell me the result.",
"The command `sleep 120` ran for 120 seconds and completed successfully.",
[]fantasy.StepResult{
stepWithToolResults(toolError("exec", "Command timed out after 8s")),
},
)
if got != "Command timed out after 8s" {
t.Fatalf("expected exec timeout grounding, got %q", got)
}
}
func stepWithToolResults(results ...fantasy.ToolResultContent) fantasy.StepResult {
content := make(fantasy.ResponseContent, 0, len(results))
for _, result := range results {
content = append(content, result)
}
return fantasy.StepResult{
Response: fantasy.Response{
Content: content,
},
}
}
func stepWithTextAndToolResults(text string, results ...fantasy.ToolResultContent) fantasy.StepResult {
content := fantasy.ResponseContent{fantasy.TextContent{Text: text}}
for _, result := range results {
content = append(content, result)
}
return fantasy.StepResult{
Response: fantasy.Response{
Content: content,
},
}
}
func toolText(toolName, text string) fantasy.ToolResultContent {
return fantasy.ToolResultContent{
ToolCallID: toolName + "-call",
ToolName: toolName,
Result: fantasy.ToolResultOutputContentText{Text: text},
}
}
func toolError(toolName, text string) fantasy.ToolResultContent {
return fantasy.ToolResultContent{
ToolCallID: toolName + "-call",
ToolName: toolName,
Result: fantasy.ToolResultOutputContentError{
Error: errors.New(text),
},
}
}

View file

@ -4,6 +4,7 @@ import (
"fmt" "fmt"
"os" "os"
"path/filepath" "path/filepath"
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg/bus" "github.com/ZanzyTHEbar/dragonscale/pkg/bus"
"github.com/ZanzyTHEbar/dragonscale/pkg/config" "github.com/ZanzyTHEbar/dragonscale/pkg/config"
@ -49,7 +50,11 @@ func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msg
registry.Register(tools.NewAppendFileTool(workspace, restrict)) registry.Register(tools.NewAppendFileTool(workspace, restrict))
// Shell execution // Shell execution
registry.Register(tools.NewExecTool(workspace, restrict)) execTool := tools.NewExecTool(workspace, restrict)
if os.Getenv("DRAGONSCALE_EVAL_CONFIG") != "" || os.Getenv("DRAGONSCALE_EVAL_RUNTIME") != "" {
execTool.SetTimeout(8 * time.Second)
}
registry.Register(execTool)
if searchTool := tools.NewWebSearchTool(tools.WebSearchToolOptions{ if searchTool := tools.NewWebSearchTool(tools.WebSearchToolOptions{
BraveAPIKey: cfg.Tools.Web.Brave.APIKey, BraveAPIKey: cfg.Tools.Web.Brave.APIKey,

View file

@ -0,0 +1,140 @@
package agent
import (
"context"
"testing"
"github.com/ZanzyTHEbar/dragonscale/pkg/skills"
"github.com/ZanzyTHEbar/dragonscale/pkg/tools"
)
func TestInitialPromptToolsExposeSkillAndFileHelpers(t *testing.T) {
reg := tools.NewToolRegistry()
reg.Register(&namedTool{name: "tool_search"})
reg.Register(&namedTool{name: "tool_call"})
reg.Register(tools.NewExecTool(t.TempDir(), true))
reg.Register(tools.NewReadFileTool(t.TempDir(), true))
reg.Register(tools.NewWriteFileTool(t.TempDir(), true))
reg.Register(tools.NewEditFileTool(t.TempDir(), true))
reg.Register(tools.NewAppendFileTool(t.TempDir(), true))
reg.Register(tools.NewListDirTool(t.TempDir(), true))
reg.Register(&namedTool{name: "spawn"})
reg.Register(&namedTool{name: "subagent"})
reg.Register(&namedTool{name: "memory"})
reg.Register(&namedTool{name: "obligation"})
skillsDir := t.TempDir()
loader := skills.NewSkillsLoader(skillsDir, "", "")
reg.Register(tools.NewSkillSearchTool(loader))
reg.Register(tools.NewSkillReadTool(loader))
reg.Register(tools.NewSkillTraverseTool(loader))
al := &AgentLoop{tools: reg}
skillNames := toolNames(al.initialPromptTools("Read the 'eval-test-skill' skill and explain it."))
if !containsAll(skillNames, "skill_search", "skill_read") {
t.Fatalf("expected skill tools, got %v", skillNames)
}
editNames := toolNames(al.initialPromptTools("Edit notes.txt to replace foo with bar."))
if !containsAll(editNames, "edit_file") {
t.Fatalf("expected edit_file, got %v", editNames)
}
appendNames := toolNames(al.initialPromptTools("Append a final line to report.txt."))
if !containsAll(appendNames, "append_file") {
t.Fatalf("expected append_file, got %v", appendNames)
}
execNames := toolNames(al.initialPromptTools("Run the command 'echo dragonscale-eval-test' and tell me the output."))
if !containsAll(execNames, "exec") {
t.Fatalf("expected exec, got %v", execNames)
}
writeReadNames := toolNames(al.initialPromptTools("Write the text 'dragonscale eval checkpoint' to a file called eval_checkpoint.txt, then read it back and confirm the contents match."))
if !containsAll(writeReadNames, "write_file", "read_file") {
t.Fatalf("expected write_file/read_file, got %v", writeReadNames)
}
listNames := toolNames(al.initialPromptTools("Create a file called project/readme.txt with 'Project initialized'. Then list the project directory to verify it exists."))
if !containsAll(listNames, "write_file", "list_dir") {
t.Fatalf("expected write_file/list_dir, got %v", listNames)
}
spawnNames := toolNames(al.initialPromptTools("Spawn a background task to write the text 'async-spawn-test' to a file called spawn_output.txt."))
if !containsAll(spawnNames, "spawn", "write_file") {
t.Fatalf("expected spawn/write_file, got %v", spawnNames)
}
subagentNames := toolNames(al.initialPromptTools("Use a subagent to calculate the sum of 10 + 20 + 30 and report the result back to me."))
if !containsAll(subagentNames, "subagent") {
t.Fatalf("expected subagent, got %v", subagentNames)
}
memoryNames := toolNames(al.initialPromptTools("Track these commitments exactly: send rent receipt tonight, book vet appointment tomorrow, and submit sprint notes by Friday."))
if !containsAll(memoryNames, "memory") {
t.Fatalf("expected memory, got %v", memoryNames)
}
memorySearchNames := toolNames(al.initialPromptTools("Search your memory for 'xyzzy_nonexistent_topic_42' and tell me what you find."))
if !containsAll(memorySearchNames, "memory") {
t.Fatalf("expected memory for memory-search prompt, got %v", memorySearchNames)
}
memoryStatusNames := toolNames(al.initialPromptTools("Check your memory system status and tell me the current context pressure level."))
if !containsAll(memoryStatusNames, "memory") {
t.Fatalf("expected memory for memory-status prompt, got %v", memoryStatusNames)
}
discoveryNames := toolNames(al.initialPromptTools("Search for a tool that can read files."))
if !containsAll(discoveryNames, "tool_search") || containsAll(discoveryNames, "read_file") {
t.Fatalf("expected discovery prompt to expose only tool_search-like helpers, got %v", discoveryNames)
}
}
func TestIsPlanningOnlyPrompt(t *testing.T) {
t.Parallel()
if !isPlanningOnlyPrompt("Create a 6-week proactive check-in schedule for learning Spanish with weekly milestones.") {
t.Fatal("expected planning-only prompt to be detected")
}
if isPlanningOnlyPrompt("Run 'date +%Y' to get the current year, write that year to a file, then read it back.") {
t.Fatal("expected action-oriented prompt not to be treated as planning-only")
}
if isPlanningOnlyPrompt("Capture these commitments and give me a reminder plan.") {
t.Fatal("expected capture/reminder prompt not to be treated as planning-only")
}
if !isPlanningOnlyPrompt("I must send a proposal in 4 hours. Give me a reminder schedule and specify when the first reminder should fire.") {
t.Fatal("expected reminder schedule request to be treated as planning-only")
}
}
func containsAll(have []string, want ...string) bool {
set := make(map[string]struct{}, len(have))
for _, name := range have {
set[name] = struct{}{}
}
for _, name := range want {
if _, ok := set[name]; !ok {
return false
}
}
return true
}
type namedTool struct {
name string
}
func (n *namedTool) Name() string { return n.name }
func (n *namedTool) Description() string { return n.name }
func (n *namedTool) Parameters() map[string]interface{} { return map[string]interface{}{} }
func (n *namedTool) Execute(_ context.Context, _ map[string]interface{}) *tools.ToolResult {
return &tools.ToolResult{ForLLM: n.name}
}

View file

@ -39,43 +39,46 @@ import (
) )
type AgentLoop struct { type AgentLoop struct {
bus *bus.MessageBus bus *bus.MessageBus
languageModel fantasy.LanguageModel languageModel fantasy.LanguageModel
workspace string workspace string
model string model string
contextWindow int // Maximum context window size in tokens contextWindow int // Maximum context window size in tokens
maxIterations int maxIterations int
sessions *session.SessionManager sessions *session.SessionManager
state *state.Manager state *state.Manager
contextBuilder *ContextBuilder contextBuilder *ContextBuilder
tools *tools.ToolRegistry activeContextBuilder *DefaultActiveContextBuilder
memoryStore *memstore.MemoryStore // 3-tier MemGPT memory (always initialized) tools *tools.ToolRegistry
memDelegate memory.MemoryDelegate // DB delegate (always initialized) memoryStore *memstore.MemoryStore // 3-tier MemGPT memory (always initialized)
obsManager *observation.Manager // Observational memory (always initialized) memDelegate memory.MemoryDelegate // DB delegate (always initialized)
secureBus *securebus.Bus // ITR SecureBus (always initialized) obsManager *observation.Manager // Observational memory (always initialized)
queries *memsqlc.Queries // SQL query surface for runtime persistence secureBus *securebus.Bus // ITR SecureBus (always initialized)
kvDelegate KVDelegate // KV adapter for offloaded tool results queries *memsqlc.Queries // SQL query surface for runtime persistence
stateStore *StateStore // Agent run state persistence kvDelegate KVDelegate // KV adapter for offloaded tool results
offloadThresholdChars int // Char threshold for tool result offloading (derived from token config) stateStore *StateStore // Agent run state persistence
conversationIDs *boundedCache[string, ids.UUID] // Owner: agent_run.go — wrote by prepareRuntimeState, read in prepareRuntimeState/load path offloadThresholdChars int // Char threshold for tool result offloading (derived from token config)
conversationMu sync.Mutex // serializes conversation creation path rlmEngine rlmAnswerer // Recursive context reducer for oversized historical segments
identitySync *dragonsync.IdentitySync // File→DB sync for identity docs (nil if memory disabled) rlmDirectThresholdBytes int // Byte threshold before invoking RLM reduction
activeSessionKey atomic.Value // Owner: agent_run.go — written in runAgentLoop, read by router/toolloop for context routing conversationIDs *boundedCache[string, ids.UUID] // Owner: agent_run.go — wrote by prepareRuntimeState, read in prepareRuntimeState/load path
running atomic.Bool // Owner: loop.go — lifecycle gate controlled by Run/Stop only conversationMu sync.Mutex // serializes conversation creation path
summarizing sync.Map // Owner: summarizer.go — intended for async summarization lockout, currently gated by TODO path identitySync *dragonsync.IdentitySync // File→DB sync for identity docs (nil if memory disabled)
summarizeFailures sync.Map // Owner: summarizer.go — write/read in forceCompression + summarizeSession error paths activeSessionKey atomic.Value // Owner: agent_run.go — written in runAgentLoop, read by router/toolloop for context routing
contextTreeCache sync.Map // Owner: summarizer.go — sessionKey → contextTreeCacheEntry keyed by query and history size running atomic.Bool // Owner: loop.go — lifecycle gate controlled by Run/Stop only
auditChan chan *memory.AuditEntry // Buffered channel for async audit logging; drained by background worker summarizing sync.Map // Owner: summarizer.go — intended for async summarization lockout, currently gated by TODO path
auditDone chan struct{} // Closed when audit worker exits summarizeFailures sync.Map // Owner: summarizer.go — write/read in forceCompression + summarizeSession error paths
focusDirty sync.Map // sessionKey → struct{}: set by focus tool callbacks, cleared after context reload contextTreeCache sync.Map // Owner: summarizer.go — sessionKey → contextTreeCacheEntry keyed by query and history size
ctxBlockCache sync.Map // sessionKey → ctxBlockCacheEntry: cached focus + knowledge blocks auditChan chan *memory.AuditEntry // Buffered channel for async audit logging; drained by background worker
cfg *config.Config // Stored for subagent factory access auditDone chan struct{} // Closed when audit worker exits
channelManager *channels.Manager focusDirty sync.Map // sessionKey → struct{}: set by focus tool callbacks, cleared after context reload
commandRegistry []SlashCommand ctxBlockCache sync.Map // sessionKey → ctxBlockCacheEntry: cached focus + knowledge blocks
outputOverride atomic.Value // Owner: command_handler.go — CLI output redirection target for internal messages cfg *config.Config // Stored for subagent factory access
toolResultSearch fantasy.AgentTool channelManager *channels.Manager
cortex *cortex.Cortex commandRegistry []SlashCommand
inflight sync.WaitGroup outputOverride atomic.Value // Owner: command_handler.go — CLI output redirection target for internal messages
toolResultSearch fantasy.AgentTool
cortex *cortex.Cortex
inflight sync.WaitGroup
} }
type outputTarget struct { type outputTarget struct {
@ -197,7 +200,8 @@ func NewAgentLoop(ctx context.Context, cfg *config.Config, msgBus *bus.MessageBu
// Register memory/search/skill tools on subagent registry so spawned // Register memory/search/skill tools on subagent registry so spawned
// agents can search knowledge, offload results, and use skills. // agents can search knowledge, offload results, and use skills.
subagentTools.Register(NewMemGPTTool(ms, pkg.NAME, "default")) subagentMemTool := NewMemGPTTool(ms, pkg.NAME, "default")
subagentTools.Register(subagentMemTool)
subagentTools.Register(tools.NewObligationTool(memDelegate, pkg.NAME)) subagentTools.Register(tools.NewObligationTool(memDelegate, pkg.NAME))
subagentTools.Register(tools.NewKeywordSearchTool(ms, pkg.NAME)) subagentTools.Register(tools.NewKeywordSearchTool(ms, pkg.NAME))
subagentTools.Register(tools.NewSemanticSearchTool(ms, pkg.NAME)) subagentTools.Register(tools.NewSemanticSearchTool(ms, pkg.NAME))
@ -290,37 +294,41 @@ func NewAgentLoop(ctx context.Context, cfg *config.Config, msgBus *bus.MessageBu
return resp.Content.Text(), nil return resp.Content.Text(), nil
} }
obsManager := observation.NewManager(memDelegate, pkg.NAME, callModelFn, observation.DefaultManagerConfig()) obsManager := observation.NewManager(memDelegate, pkg.NAME, callModelFn, observation.DefaultManagerConfig())
rlmEngine, rlmThresholdBytes := newLiveRLMAnswerer(model)
auditCh := make(chan *memory.AuditEntry, 256) auditCh := make(chan *memory.AuditEntry, 256)
auditDone := make(chan struct{}) auditDone := make(chan struct{})
al := &AgentLoop{ al := &AgentLoop{
bus: msgBus, bus: msgBus,
languageModel: model, languageModel: model,
workspace: workspace, workspace: workspace,
model: cfg.Agents.Defaults.Model, model: cfg.Agents.Defaults.Model,
contextWindow: cfg.Agents.Defaults.MaxTokens, contextWindow: cfg.Agents.Defaults.MaxTokens,
maxIterations: cfg.Agents.Defaults.MaxToolIterations, maxIterations: cfg.Agents.Defaults.MaxToolIterations,
sessions: sessionsManager, sessions: sessionsManager,
state: stateManager, state: stateManager,
contextBuilder: contextBuilder, contextBuilder: contextBuilder,
tools: toolsRegistry, tools: toolsRegistry,
memoryStore: ms, memoryStore: ms,
memDelegate: memDelegate, memDelegate: memDelegate,
obsManager: obsManager, obsManager: obsManager,
queries: queries, queries: queries,
kvDelegate: kv, kvDelegate: kv,
stateStore: stateStore, stateStore: stateStore,
offloadThresholdChars: offloadThreshold * 4, offloadThresholdChars: offloadThreshold * 4,
toolResultSearch: NewToolResultSearchTool(queries, kv), rlmEngine: rlmEngine,
conversationIDs: newBoundedCache[string, ids.UUID](1024), rlmDirectThresholdBytes: rlmThresholdBytes,
identitySync: idSync, toolResultSearch: NewToolResultSearchTool(queries, kv),
summarizing: sync.Map{}, conversationIDs: newBoundedCache[string, ids.UUID](1024),
auditChan: auditCh, identitySync: idSync,
auditDone: auditDone, summarizing: sync.Map{},
commandRegistry: defaultSlashCommands(), auditChan: auditCh,
cfg: cfg, auditDone: auditDone,
commandRegistry: defaultSlashCommands(),
cfg: cfg,
} }
al.activeContextBuilder = NewDefaultActiveContextBuilder(pkg.NAME, contextBuilder, sessionsManager, memDelegate, ms, queries)
go al.auditWorker(ctx, auditCh, auditDone) go al.auditWorker(ctx, auditCh, auditDone)
@ -337,6 +345,9 @@ func NewAgentLoop(ctx context.Context, cfg *config.Config, msgBus *bus.MessageBu
} }
return "" return ""
} }
contextBuilder.SetSessionResolver(sessionKeyFn)
memTool.SetSessionResolver(sessionKeyFn)
subagentMemTool.SetSessionResolver(sessionKeyFn)
focusInvalidate := func() { focusInvalidate := func() {
if sk := sessionKeyFn(); sk != "" { if sk := sessionKeyFn(); sk != "" {
al.focusDirty.Store(sk, struct{}{}) al.focusDirty.Store(sk, struct{}{})

View file

@ -7,6 +7,8 @@ package agent
import ( import (
"context" "context"
"strings"
jsonv2 "github.com/go-json-experiment/json" jsonv2 "github.com/go-json-experiment/json"
memstore "github.com/ZanzyTHEbar/dragonscale/pkg/memory/store" memstore "github.com/ZanzyTHEbar/dragonscale/pkg/memory/store"
@ -16,7 +18,10 @@ import (
// MemGPTTool wraps store.MemoryTool as a DragonScale tools.Tool so it can be // MemGPTTool wraps store.MemoryTool as a DragonScale tools.Tool so it can be
// registered in the ToolRegistry and executed by the Fantasy agent loop. // registered in the ToolRegistry and executed by the Fantasy agent loop.
type MemGPTTool struct { type MemGPTTool struct {
inner *memstore.MemoryTool store *memstore.MemoryStore
agentID string
session string
sessionKeyFn func() string
} }
var _ tools.Tool = (*MemGPTTool)(nil) var _ tools.Tool = (*MemGPTTool)(nil)
@ -24,7 +29,9 @@ var _ tools.Tool = (*MemGPTTool)(nil)
// NewMemGPTTool creates a DragonScale tool wrapper around a MemoryTool. // NewMemGPTTool creates a DragonScale tool wrapper around a MemoryTool.
func NewMemGPTTool(store *memstore.MemoryStore, agentID, session string) *MemGPTTool { func NewMemGPTTool(store *memstore.MemoryStore, agentID, session string) *MemGPTTool {
return &MemGPTTool{ return &MemGPTTool{
inner: memstore.NewMemoryTool(store, agentID, session), store: store,
agentID: agentID,
session: session,
} }
} }
@ -89,7 +96,7 @@ func (t *MemGPTTool) Execute(ctx context.Context, args map[string]interface{}) *
return tools.ErrorResult("invalid arguments: " + err.Error()) return tools.ErrorResult("invalid arguments: " + err.Error())
} }
result, err := t.inner.Execute(ctx, string(input)) result, err := memstore.NewMemoryTool(t.store, t.agentID, t.currentSession()).Execute(ctx, string(input))
if err != nil { if err != nil {
return tools.ErrorResult("memory tool error: " + err.Error()) return tools.ErrorResult("memory tool error: " + err.Error())
} }
@ -102,5 +109,23 @@ func (t *MemGPTTool) Execute(ctx context.Context, args map[string]interface{}) *
// UpdateSession rebinds the inner MemoryTool to a new session. // UpdateSession rebinds the inner MemoryTool to a new session.
// Called when the agent switches sessions. // Called when the agent switches sessions.
func (t *MemGPTTool) UpdateSession(store *memstore.MemoryStore, agentID, session string) { func (t *MemGPTTool) UpdateSession(store *memstore.MemoryStore, agentID, session string) {
t.inner = memstore.NewMemoryTool(store, agentID, session) t.store = store
t.agentID = agentID
t.session = session
}
func (t *MemGPTTool) SetSessionResolver(sessionKeyFn func() string) {
t.sessionKeyFn = sessionKeyFn
}
func (t *MemGPTTool) currentSession() string {
if t.sessionKeyFn != nil {
if sessionKey := strings.TrimSpace(t.sessionKeyFn()); sessionKey != "" {
return sessionKey
}
}
if strings.TrimSpace(t.session) != "" {
return t.session
}
return "default"
} }

View file

@ -16,23 +16,6 @@ import (
const defaultToolMaxConcurrency = 4 const defaultToolMaxConcurrency = 4
type ctxStepIndexKey struct{}
func WithStepIndex(ctx context.Context, stepIndex int) context.Context {
return context.WithValue(ctx, ctxStepIndexKey{}, stepIndex)
}
func StepIndexFromCtx(ctx context.Context) int {
v := ctx.Value(ctxStepIndexKey{})
if v == nil {
return 0
}
if i, ok := v.(int); ok {
return i
}
return 0
}
// OffloadingToolRuntime wraps a base ToolRuntime and applies tool result // OffloadingToolRuntime wraps a base ToolRuntime and applies tool result
// offloading policy: // offloading policy:
// - Always offload full results to KV delegate. // - Always offload full results to KV delegate.
@ -78,7 +61,7 @@ func (r OffloadingToolRuntime) Execute(ctx context.Context, tools []fantasy.Agen
chunkChars = 2_000 chunkChars = 2_000
} }
stepIndex := StepIndexFromCtx(ctx) stepIndex := fantasy.StepIndexFromCtx(ctx)
results, err := r.Base.Execute(ctx, tools, toolCalls, nil) results, err := r.Base.Execute(ctx, tools, toolCalls, nil)
if err != nil { if err != nil {

169
pkg/agent/rlm_runtime.go Normal file
View file

@ -0,0 +1,169 @@
package agent
import (
"context"
"fmt"
"strings"
fantasy "charm.land/fantasy"
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/observation"
"github.com/ZanzyTHEbar/dragonscale/pkg/rlm"
)
type rlmAnswerer interface {
Answer(ctx context.Context, sessionKey, query, context_ string) (string, uint32, error)
}
const rlmReducerSource = "rlm_context_reducer"
func newLiveRLMAnswerer(model fantasy.LanguageModel) (rlmAnswerer, int) {
cfg := rlm.DefaultEngineConfig()
engine := rlm.NewEngine(cfg, nil, makeRLMCallModel(model))
return engine, cfg.Strategy.DirectThreshold
}
func makeRLMCallModel(model fantasy.LanguageModel) rlm.CallModelFunc {
return func(ctx context.Context, ctxContent, query string) (string, uint32, error) {
systemPrompt := "You reduce large runtime context for an autonomous agent. Extract only the facts, constraints, prior decisions, and tool outcomes that are relevant to the user query. Return a compact context note, not a final assistant reply."
userPrompt := fmt.Sprintf("Query:\n%s\n\nContext:\n%s", query, ctxContent)
if strings.HasPrefix(ctxContent, "You are synthesising answers from multiple text partitions.") {
systemPrompt = ctxContent
userPrompt = query
}
maxTokens := int64(384)
resp, err := model.Generate(ctx, fantasy.Call{
Prompt: fantasy.Prompt{
fantasy.NewSystemMessage(systemPrompt),
fantasy.NewUserMessage(userPrompt),
},
MaxOutputTokens: &maxTokens,
})
if err != nil {
return "", 0, err
}
return strings.TrimSpace(resp.Content.Text()), uint32(resp.Usage.TotalTokens), nil
}
}
func (al *AgentLoop) maybeReduceProjectionWithRLM(ctx context.Context, sessionKey, query string, projection *memory.ActiveContextProjection) *memory.ActiveContextProjection {
if al == nil || al.rlmEngine == nil || projection == nil {
return projection
}
candidateSegments := make([]memory.ProjectionSegment, 0, len(projection.Segments))
for _, seg := range projection.Segments {
if isRLMReducibleSegment(seg.Kind) {
candidateSegments = append(candidateSegments, seg)
}
}
if len(candidateSegments) == 0 {
return projection
}
contextBlob := renderProjectionSegmentsForRLM(candidateSegments)
if strings.TrimSpace(contextBlob) == "" || len(contextBlob) < al.rlmDirectThresholdBytes {
return projection
}
reduced, _, err := al.rlmEngine.Answer(ctx, sessionKey, query, contextBlob)
if err != nil {
logger.WarnCF("agent", "RLM context reduction failed", map[string]any{
"session_key": sessionKey,
"error": err.Error(),
})
return projection
}
reduced = strings.TrimSpace(reduced)
if reduced == "" {
return projection
}
rlmSegment := memory.ProjectionSegment{
Kind: memory.ProjectionSegmentSystem,
Source: rlmReducerSource,
Text: "## Recursive Context Reduction\n\n" + reduced,
Tokens: observation.EstimateTokens(reduced),
Ref: mergeProjectionRefs(sessionKey, candidateSegments),
}
updated := make([]memory.ProjectionSegment, 0, len(projection.Segments)-len(candidateSegments)+1)
inserted := false
for _, seg := range projection.Segments {
if isRLMReducibleSegment(seg.Kind) {
if !inserted {
updated = append(updated, rlmSegment)
inserted = true
}
continue
}
updated = append(updated, seg)
}
clone := *projection
clone.Segments = updated
return &clone
}
func isRLMReducibleSegment(kind memory.ProjectionSegmentKind) bool {
switch kind {
case memory.ProjectionSegmentDAG, memory.ProjectionSegmentRecall, memory.ProjectionSegmentArchival:
return true
default:
return false
}
}
func renderProjectionSegmentsForRLM(segments []memory.ProjectionSegment) string {
var sb strings.Builder
for _, seg := range segments {
text := strings.TrimSpace(seg.Text)
if text == "" {
continue
}
sb.WriteString("### ")
sb.WriteString(string(seg.Kind))
if seg.Source != "" {
sb.WriteString(" / ")
sb.WriteString(seg.Source)
}
sb.WriteString("\n")
sb.WriteString(text)
sb.WriteString("\n\n")
}
return strings.TrimSpace(sb.String())
}
func mergeProjectionRefs(sessionKey string, segments []memory.ProjectionSegment) memory.ImmutableSpanRef {
var merged memory.ImmutableSpanRef
for _, seg := range segments {
ref := seg.Ref
if ref.SessionKey == "" {
continue
}
if merged.SessionKey == "" {
merged = ref
continue
}
if ref.StartIdx < merged.StartIdx {
merged.StartIdx = ref.StartIdx
merged.FirstID = ref.FirstID
}
if merged.FromTime.IsZero() || (!ref.FromTime.IsZero() && ref.FromTime.Before(merged.FromTime)) {
merged.FromTime = ref.FromTime
}
if ref.EndIdx > merged.EndIdx {
merged.EndIdx = ref.EndIdx
merged.LastID = ref.LastID
}
if merged.ToTime.IsZero() || ref.ToTime.After(merged.ToTime) {
merged.ToTime = ref.ToTime
}
}
if merged.SessionKey == "" {
return memory.ImmutableSpanRef{SessionKey: sessionKey}
}
return merged
}

View file

@ -0,0 +1,114 @@
package agent
import (
"context"
"os"
"path/filepath"
"testing"
pkgroot "github.com/ZanzyTHEbar/dragonscale/pkg"
"github.com/ZanzyTHEbar/dragonscale/pkg/bus"
"github.com/ZanzyTHEbar/dragonscale/pkg/config"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory/dag"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAssembleContext_ReducesProjectionWithRLM(t *testing.T) {
t.Parallel()
tmpDir, err := os.MkdirTemp("", "agent-rlm-*")
require.NoError(t, err)
defer os.RemoveAll(tmpDir)
cfg := config.DefaultConfig()
cfg.Agents.Defaults.Sandbox = tmpDir
cfg.Agents.Defaults.MaxTokens = 4096
cfg.Memory.DBPath = filepath.Join(tmpDir, "agent-rlm.db")
al := mustNewAgentLoop(t, cfg, bus.NewMessageBus(), newMockLanguageModel("ok"))
mock := &mockRLMAnswerer{answer: "critical dependency: prior DAG summary says to preserve the write-path invariants"}
al.rlmEngine = mock
al.rlmDirectThresholdBytes = 1
sessionKey := "rlm-session"
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "user",
Content: "one",
TokenEstimate: 4,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "assistant",
Content: "two",
TokenEstimate: 4,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "user",
Content: "three",
TokenEstimate: 4,
})
insertImmutableMessage(t, al, &memory.ImmutableMessage{
ID: ids.New(),
SessionKey: sessionKey,
Role: "assistant",
Content: "four",
TokenEstimate: 4,
})
persister, ok := al.memDelegate.(dag.DAGPersister)
require.True(t, ok)
tree := dag.NewDAG()
tree.Add(&dag.Node{
ID: "session-root",
Level: dag.LevelSession,
Summary: "DAG summary: preserve the write-path invariants and downstream dependency ordering",
Tokens: 32,
StartIdx: 0,
EndIdx: 4,
})
tree.SetRoots([]string{"session-root"})
require.NoError(t, persister.PersistDAG(t.Context(), pkgroot.NAME, sessionKey, &dag.PersistSnapshot{
FromMsgIdx: 0,
ToMsgIdx: 4,
MsgCount: 4,
DAG: tree,
}))
ac, err := al.assembleContext(t.Context(), processOptions{
SessionKey: sessionKey,
UserMessage: "What dependency matters?",
EnableSummary: false,
})
require.NoError(t, err)
require.NotNil(t, ac.projection)
assert.Equal(t, 1, mock.calls)
assert.Equal(t, "What dependency matters?", mock.query)
assert.Contains(t, mock.context, "DAG summary: preserve the write-path invariants")
assert.Contains(t, ac.systemPrompt, "Recursive Context Reduction")
assert.Contains(t, ac.systemPrompt, mock.answer)
assert.NotContains(t, projectionKinds(ac.projection), memory.ProjectionSegmentDAG)
}
type mockRLMAnswerer struct {
answer string
context string
query string
calls int
}
func (m *mockRLMAnswerer) Answer(_ context.Context, _ string, query, context_ string) (string, uint32, error) {
m.calls++
m.query = query
m.context = context_
return m.answer, 17, nil
}

View file

@ -0,0 +1,311 @@
package agent
import (
"context"
"fmt"
"os"
"testing"
"time"
fantasy "charm.land/fantasy"
pkgroot "github.com/ZanzyTHEbar/dragonscale/pkg"
"github.com/ZanzyTHEbar/dragonscale/pkg/bus"
"github.com/ZanzyTHEbar/dragonscale/pkg/config"
memsqlc "github.com/ZanzyTHEbar/dragonscale/pkg/memory/sqlc"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type multiToolCallingModel struct{}
func (m *multiToolCallingModel) Generate(_ context.Context, call fantasy.Call) (*fantasy.Response, error) {
toolResults := countPromptToolResults(call.Prompt)
switch toolResults {
case 0:
return &fantasy.Response{
Content: fantasy.ResponseContent{
fantasy.ToolCallContent{ToolCallID: "call-1", ToolName: "echo", Input: `{"text":"one"}`},
},
FinishReason: fantasy.FinishReasonToolCalls,
Usage: fantasy.Usage{InputTokens: 5, OutputTokens: 3, TotalTokens: 8},
}, nil
case 1:
return &fantasy.Response{
Content: fantasy.ResponseContent{
fantasy.ToolCallContent{ToolCallID: "call-2", ToolName: "echo", Input: `{"text":"two"}`},
},
FinishReason: fantasy.FinishReasonToolCalls,
Usage: fantasy.Usage{InputTokens: 6, OutputTokens: 3, TotalTokens: 9},
}, nil
default:
return &fantasy.Response{
Content: fantasy.ResponseContent{fantasy.TextContent{Text: "Final response after two tools"}},
FinishReason: fantasy.FinishReasonStop,
Usage: fantasy.Usage{InputTokens: 7, OutputTokens: 5, TotalTokens: 12},
}, nil
}
}
func (m *multiToolCallingModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) {
resp, err := m.Generate(ctx, call)
if err != nil {
return nil, err
}
return func(yield func(fantasy.StreamPart) bool) {
if len(resp.Content.ToolCalls()) > 0 {
for _, tc := range resp.Content.ToolCalls() {
if !yield(fantasy.StreamPart{
Type: fantasy.StreamPartTypeToolCall,
ID: tc.ToolCallID,
ToolCallName: tc.ToolName,
ToolCallInput: tc.Input,
}) {
return
}
}
yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls, Usage: resp.Usage})
return
}
text := resp.Content.Text()
if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "text-0"}) {
return
}
if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "text-0", Delta: text}) {
return
}
if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "text-0"}) {
return
}
yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop, Usage: resp.Usage})
}, nil
}
func (m *multiToolCallingModel) GenerateObject(_ context.Context, _ fantasy.ObjectCall) (*fantasy.ObjectResponse, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *multiToolCallingModel) StreamObject(_ context.Context, _ fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *multiToolCallingModel) Provider() string { return "mock" }
func (m *multiToolCallingModel) Model() string { return "multi-tool-model" }
type sameStepMultiToolModel struct{}
func (m *sameStepMultiToolModel) Generate(_ context.Context, call fantasy.Call) (*fantasy.Response, error) {
if countPromptToolResults(call.Prompt) == 0 {
return &fantasy.Response{
Content: fantasy.ResponseContent{
fantasy.ToolCallContent{ToolCallID: "call-1", ToolName: "echo", Input: `{"text":"one"}`},
fantasy.ToolCallContent{ToolCallID: "call-2", ToolName: "echo", Input: `{"text":"two"}`},
},
FinishReason: fantasy.FinishReasonToolCalls,
Usage: fantasy.Usage{InputTokens: 5, OutputTokens: 4, TotalTokens: 9},
}, nil
}
return &fantasy.Response{
Content: fantasy.ResponseContent{fantasy.TextContent{Text: "Final response after parallel tools"}},
FinishReason: fantasy.FinishReasonStop,
Usage: fantasy.Usage{InputTokens: 6, OutputTokens: 5, TotalTokens: 11},
}, nil
}
func (m *sameStepMultiToolModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) {
resp, err := m.Generate(ctx, call)
if err != nil {
return nil, err
}
return func(yield func(fantasy.StreamPart) bool) {
if len(resp.Content.ToolCalls()) > 0 {
for _, tc := range resp.Content.ToolCalls() {
if !yield(fantasy.StreamPart{
Type: fantasy.StreamPartTypeToolCall,
ID: tc.ToolCallID,
ToolCallName: tc.ToolName,
ToolCallInput: tc.Input,
}) {
return
}
}
yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls, Usage: resp.Usage})
return
}
text := resp.Content.Text()
if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "text-0"}) {
return
}
if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "text-0", Delta: text}) {
return
}
if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "text-0"}) {
return
}
yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop, Usage: resp.Usage})
}, nil
}
func (m *sameStepMultiToolModel) GenerateObject(_ context.Context, _ fantasy.ObjectCall) (*fantasy.ObjectResponse, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *sameStepMultiToolModel) StreamObject(_ context.Context, _ fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) {
return nil, fmt.Errorf("not implemented")
}
func (m *sameStepMultiToolModel) Provider() string { return "mock" }
func (m *sameStepMultiToolModel) Model() string { return "same-step-multi-tool-model" }
func countPromptToolResults(prompt []fantasy.Message) int {
count := 0
for _, msg := range prompt {
for _, part := range msg.Content {
if part.GetType() == fantasy.ContentTypeToolResult {
count++
}
}
}
return count
}
func uniqueTransitionSteps(rows []memsqlc.AgentStateTransition) []int64 {
seen := make(map[int64]struct{})
steps := make([]int64, 0, len(rows))
for _, row := range rows {
if _, ok := seen[row.StepIndex]; ok {
continue
}
seen[row.StepIndex] = struct{}{}
steps = append(steps, row.StepIndex)
}
return steps
}
func TestIntegration_RuntimeBookkeeping_PersistsTransitionsAndMetrics(t *testing.T) {
t.Parallel()
tmpDir, err := os.MkdirTemp("", "agent-runtime-bookkeeping-*")
require.NoError(t, err)
defer os.RemoveAll(tmpDir)
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Sandbox: tmpDir,
Model: "multi-tool-model",
MaxTokens: 4096,
MaxToolIterations: 10,
},
},
}
msgBus := bus.NewMessageBus()
al := mustNewAgentLoop(t, cfg, msgBus, &multiToolCallingModel{})
al.RegisterTool(&echoTool{})
ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
defer cancel()
since := time.Now().Add(-time.Second)
msg := bus.InboundMessage{
Channel: "test",
SenderID: "user1",
ChatID: "chat1",
Content: "Use the echo tool twice",
SessionKey: "runtime-bookkeeping-session",
}
response, err := al.processMessage(ctx, msg)
require.NoError(t, err)
assert.Contains(t, response, "Final response after two tools")
completions, err := al.queries.GetCompletedTasks(ctx, memsqlc.GetCompletedTasksParams{
AgentID: pkgroot.NAME,
Since: since,
})
require.NoError(t, err)
require.NotEmpty(t, completions)
completion := completions[len(completions)-1]
require.NotNil(t, completion.ToolCalls)
assert.Equal(t, int64(2), *completion.ToolCalls)
transitions, err := al.queries.ListAgentStateTransitionsByRunID(ctx, memsqlc.ListAgentStateTransitionsByRunIDParams{
RunID: completion.RunID,
Lim: 128,
})
require.NoError(t, err)
assert.NotEmpty(t, transitions)
assert.Equal(t, []int64{0, 1, 2}, uniqueTransitionSteps(transitions))
toolResults, err := al.queries.ListAgentToolResultsByRunID(ctx, memsqlc.ListAgentToolResultsByRunIDParams{
RunID: completion.RunID,
Lim: 16,
})
require.NoError(t, err)
require.Len(t, toolResults, 2)
assert.Equal(t, int64(0), toolResults[0].StepIndex)
assert.Equal(t, int64(1), toolResults[1].StepIndex)
}
func TestIntegration_RuntimeBookkeeping_UsesAgentStepForMultipleToolCalls(t *testing.T) {
t.Parallel()
tmpDir, err := os.MkdirTemp("", "agent-runtime-multicall-*")
require.NoError(t, err)
defer os.RemoveAll(tmpDir)
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Sandbox: tmpDir,
Model: "same-step-multi-tool-model",
MaxTokens: 4096,
MaxToolIterations: 10,
},
},
}
msgBus := bus.NewMessageBus()
al := mustNewAgentLoop(t, cfg, msgBus, &sameStepMultiToolModel{})
al.RegisterTool(&echoTool{})
ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
defer cancel()
since := time.Now().Add(-time.Second)
msg := bus.InboundMessage{
Channel: "test",
SenderID: "user1",
ChatID: "chat1",
Content: "Use two echo tool calls in one step",
SessionKey: "runtime-bookkeeping-multicall",
}
response, err := al.processMessage(ctx, msg)
require.NoError(t, err)
assert.Contains(t, response, "Final response after parallel tools")
completions, err := al.queries.GetCompletedTasks(ctx, memsqlc.GetCompletedTasksParams{
AgentID: pkgroot.NAME,
Since: since,
})
require.NoError(t, err)
require.NotEmpty(t, completions)
completion := completions[len(completions)-1]
require.NotNil(t, completion.ToolCalls)
assert.Equal(t, int64(2), *completion.ToolCalls)
toolResults, err := al.queries.ListAgentToolResultsByRunID(ctx, memsqlc.ListAgentToolResultsByRunIDParams{
RunID: completion.RunID,
Lim: 16,
})
require.NoError(t, err)
require.Len(t, toolResults, 2)
assert.Equal(t, int64(0), toolResults[0].StepIndex)
assert.Equal(t, int64(0), toolResults[1].StepIndex)
}

View file

@ -2,6 +2,7 @@ package agent
import ( import (
"context" "context"
"encoding/json"
"errors" "errors"
"fmt" "fmt"
"strings" "strings"
@ -9,6 +10,7 @@ import (
fantasy "charm.land/fantasy" fantasy "charm.land/fantasy"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids" "github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/itr" "github.com/ZanzyTHEbar/dragonscale/pkg/itr"
"github.com/ZanzyTHEbar/dragonscale/pkg/logger"
"github.com/ZanzyTHEbar/dragonscale/pkg/security/securebus" "github.com/ZanzyTHEbar/dragonscale/pkg/security/securebus"
) )
@ -31,10 +33,13 @@ type SecureBusToolRuntime struct {
// SessionKey is forwarded to bus requests for audit tracing. // SessionKey is forwarded to bus requests for audit tracing.
SessionKey string SessionKey string
// UserPrompt allows lightweight repair of placeholder tool arguments when
// the user provided an explicit literal value in the request.
UserPrompt string
// Optional state persistence for runtime execution. // Optional state persistence for runtime execution.
StateStore *StateStore StateStore *StateStore
RunID ids.UUID RunID ids.UUID
StepIndex int
} }
// Execute implements fantasy.ToolRuntime. // Execute implements fantasy.ToolRuntime.
@ -55,6 +60,7 @@ func (r SecureBusToolRuntime) Execute(
} }
results := make([]fantasy.ToolResultContent, 0, len(toolCalls)) results := make([]fantasy.ToolResultContent, 0, len(toolCalls))
stepIndex := fantasy.StepIndexFromCtx(ctx)
type deferredState struct { type deferredState struct {
step int step int
@ -64,9 +70,11 @@ func (r SecureBusToolRuntime) Execute(
var pendingStates []deferredState var pendingStates []deferredState
for i, tc := range toolCalls { for i, tc := range toolCalls {
step := r.StepIndex + i tc = repairToolCallInput(tc, r.UserPrompt)
step := stepIndex
pendingStates = append(pendingStates, deferredState{step, "tool_call", map[string]any{ pendingStates = append(pendingStates, deferredState{step, "tool_call", map[string]any{
"tool_name": tc.ToolName, "tool_name": tc.ToolName,
"tool_call_index": i,
}}) }})
reqID := ids.New().String() reqID := ids.New().String()
@ -95,7 +103,7 @@ func (r SecureBusToolRuntime) Execute(
} }
// Execute via Base runtime for the single tool call. // Execute via Base runtime for the single tool call.
baseResults, err := r.Base.Execute(ctx, tools, []fantasy.ToolCallContent{tc}, nil) baseResults, err := r.Base.Execute(fantasy.WithStepIndex(ctx, stepIndex), tools, []fantasy.ToolCallContent{tc}, nil)
if err != nil { if err != nil {
return results, err return results, err
} }
@ -106,7 +114,8 @@ func (r SecureBusToolRuntime) Execute(
} }
results = append(results, br) results = append(results, br)
pendingStates = append(pendingStates, deferredState{step, "tool_result", map[string]any{ pendingStates = append(pendingStates, deferredState{step, "tool_result", map[string]any{
"tool_name": tc.ToolName, "tool_name": tc.ToolName,
"tool_call_index": i,
}}) }})
if onResult != nil { if onResult != nil {
if err := onResult(br); err != nil { if err := onResult(br); err != nil {
@ -128,24 +137,45 @@ func (r SecureBusToolRuntime) recordRunState(ctx context.Context, stepIndex int,
if r.StateStore == nil || r.RunID.IsZero() { if r.StateStore == nil || r.RunID.IsZero() {
return return
} }
_, _ = r.StateStore.AddRunState(ctx, r.RunID, stepIndex, fantasy.ReActState(state), snapshot) if _, err := r.StateStore.AddRunState(ctx, r.RunID, stepIndex, fantasy.ReActState(state), snapshot); err != nil {
logger.WarnCF("agent", "Failed to record runtime step state", map[string]any{
"run_id": r.RunID.String(),
"step_index": stepIndex,
"state": state,
"error": err.Error(),
})
}
} }
// sanitizePolicyError strips internal details from policy/bus errors before // sanitizePolicyError strips internal details from policy/bus errors before
// they reach the LLM. The full error is preserved in audit state only. // they reach the LLM. The full error is preserved in audit state only.
func sanitizePolicyError(raw string) string { func sanitizePolicyError(raw string) string {
lower := strings.ToLower(raw)
switch { switch {
case strings.Contains(raw, "recursion depth"): case strings.Contains(lower, "command timed out"),
strings.Contains(lower, "timeout"),
strings.Contains(lower, "exceeds timeout budget"),
strings.Contains(lower, "command is required"),
strings.Contains(lower, "command cannot be empty"),
strings.Contains(lower, "no-op placeholder"),
strings.Contains(lower, "shell execution is disabled"),
strings.Contains(lower, "working_dir blocked"),
strings.Contains(lower, "path is required"),
strings.Contains(lower, "file not found"),
strings.Contains(lower, "skill") && strings.Contains(lower, "not found"),
strings.Contains(lower, "exit code"):
return raw
case strings.Contains(lower, "recursion depth"):
return "policy violation: recursion limit exceeded" return "policy violation: recursion limit exceeded"
case strings.Contains(raw, "network access denied"): case strings.Contains(lower, "network access denied"):
return "policy violation: network access denied" return "policy violation: network access denied"
case strings.Contains(raw, "filesystem access denied"): case strings.Contains(lower, "filesystem access denied"):
return "policy violation: filesystem access denied" return "policy violation: filesystem access denied"
case strings.Contains(raw, "secret injection failed"): case strings.Contains(lower, "secret injection failed"):
return "policy violation: unable to resolve required secrets" return "policy violation: unable to resolve required secrets"
case strings.Contains(raw, "invalid args JSON"): case strings.Contains(lower, "invalid args json"):
return "policy violation: invalid tool arguments" return "policy violation: invalid tool arguments"
case strings.Contains(raw, "policy violation"): case strings.Contains(lower, "policy violation"):
return "policy violation: access denied" return "policy violation: access denied"
default: default:
return "tool execution denied" return "tool execution denied"
@ -160,3 +190,209 @@ func overrideResultText(tr fantasy.ToolResultContent, text string) fantasy.ToolR
} }
return tr return tr
} }
func repairToolCallInput(tc fantasy.ToolCallContent, userPrompt string) fantasy.ToolCallContent {
if strings.TrimSpace(userPrompt) == "" || strings.TrimSpace(tc.Input) == "" {
return tc
}
switch tc.ToolName {
case "exec":
if command := explicitExecCommand(userPrompt); command != "" {
if repaired, ok := forceDirectArg(tc.Input, "command", command); ok {
tc.Input = repaired
}
}
case "skill_read":
if skillName := explicitSkillName(userPrompt); skillName != "" {
if repaired, ok := repairDirectArg(tc.Input, "name", skillName); ok {
tc.Input = repaired
}
}
case "read_file":
if path := explicitReadFilePath(userPrompt); path != "" {
if repaired, ok := repairDirectArg(tc.Input, "path", path); ok {
tc.Input = repaired
}
}
case "write_file":
if path, content := explicitWriteFileRequest(userPrompt); path != "" || content != "" {
if repaired, ok := repairWriteFileInput(tc.Input, path, content); ok {
tc.Input = repaired
}
}
case "list_dir":
if path := explicitListDirPath(userPrompt); path != "" {
if repaired, ok := repairDirectArg(tc.Input, "path", path); ok {
tc.Input = repaired
}
}
case "web_fetch":
if url := explicitFetchURL(userPrompt); url != "" {
var args map[string]any
if err := json.Unmarshal([]byte(tc.Input), &args); err == nil && isMalformedURLValue(stringArg(args["url"])) {
args["url"] = url
if encoded, err := json.Marshal(args); err == nil {
tc.Input = string(encoded)
}
}
}
case "tool_call":
if repaired, ok := repairNestedToolCallInput(tc.Input, userPrompt); ok {
tc.Input = repaired
}
}
return tc
}
func repairDirectArg(input, field, replacement string) (string, bool) {
var args map[string]any
if err := json.Unmarshal([]byte(input), &args); err != nil {
return "", false
}
current, _ := args[field].(string)
if !isPlaceholderValue(current) {
return "", false
}
args[field] = replacement
encoded, err := json.Marshal(args)
if err != nil {
return "", false
}
return string(encoded), true
}
func forceDirectArg(input, field, replacement string) (string, bool) {
var args map[string]any
if err := json.Unmarshal([]byte(input), &args); err != nil {
return "", false
}
if stringArg(args[field]) == replacement {
return "", false
}
args[field] = replacement
encoded, err := json.Marshal(args)
if err != nil {
return "", false
}
return string(encoded), true
}
func repairWriteFileInput(input, path, content string) (string, bool) {
var args map[string]any
if err := json.Unmarshal([]byte(input), &args); err != nil {
return "", false
}
changed := false
if path != "" && isPlaceholderValue(stringArg(args["path"])) {
args["path"] = path
changed = true
}
if content != "" && isPlaceholderValue(stringArg(args["content"])) {
args["content"] = content
changed = true
}
if !changed {
return "", false
}
encoded, err := json.Marshal(args)
if err != nil {
return "", false
}
return string(encoded), true
}
func repairNestedToolCallInput(input, userPrompt string) (string, bool) {
var args map[string]any
if err := json.Unmarshal([]byte(input), &args); err != nil {
return "", false
}
toolName, _ := args["tool_name"].(string)
nested, _ := args["arguments"].(map[string]any)
if nested == nil {
return "", false
}
changed := false
switch toolName {
case "exec":
if command := explicitExecCommand(userPrompt); command != "" && stringArg(nested["command"]) != command {
nested["command"] = command
changed = true
}
case "skill_read":
if skillName := explicitSkillName(userPrompt); skillName != "" && isPlaceholderValue(stringArg(nested["name"])) {
nested["name"] = skillName
changed = true
}
case "read_file":
if path := explicitReadFilePath(userPrompt); path != "" && isPlaceholderValue(stringArg(nested["path"])) {
nested["path"] = path
changed = true
}
case "write_file":
if path, content := explicitWriteFileRequest(userPrompt); path != "" || content != "" {
if path != "" && isPlaceholderValue(stringArg(nested["path"])) {
nested["path"] = path
changed = true
}
if content != "" && isPlaceholderValue(stringArg(nested["content"])) {
nested["content"] = content
changed = true
}
}
case "list_dir":
if path := explicitListDirPath(userPrompt); path != "" && isPlaceholderValue(stringArg(nested["path"])) {
nested["path"] = path
changed = true
}
case "web_fetch":
if url := explicitFetchURL(userPrompt); url != "" && isMalformedURLValue(stringArg(nested["url"])) {
nested["url"] = url
changed = true
}
}
if !changed {
return "", false
}
args["arguments"] = nested
encoded, err := json.Marshal(args)
if err != nil {
return "", false
}
return string(encoded), true
}
func stringArg(v any) string {
s, _ := v.(string)
return s
}
func isPlaceholderValue(value string) bool {
trimmed := strings.TrimSpace(value)
if trimmed == "" {
return true
}
if trimmed == ":" {
return true
}
return strings.Trim(trimmed, "{}[]() \t\r\n:;,.'\"`") == ""
}
func isMalformedURLValue(value string) bool {
trimmed := strings.TrimSpace(value)
if isPlaceholderValue(trimmed) {
return true
}
lower := strings.ToLower(trimmed)
return !strings.HasPrefix(lower, "http://") && !strings.HasPrefix(lower, "https://")
}

View file

@ -0,0 +1,133 @@
package agent
import (
"strings"
"testing"
fantasy "charm.land/fantasy"
)
func TestRepairToolCallInputRepairsDirectExecPlaceholder(t *testing.T) {
t.Parallel()
got := repairToolCallInput(
fantasy.ToolCallContent{
ToolCallID: "call-1",
ToolName: "exec",
Input: `{"command":":"}`,
},
"Run the command 'echo progressive-test-marker' and tell me the output.",
)
if !strings.Contains(got.Input, `"echo progressive-test-marker"`) {
t.Fatalf("expected repaired exec command, got %q", got.Input)
}
}
func TestRepairToolCallInputForcesExplicitExecCommand(t *testing.T) {
t.Parallel()
got := repairToolCallInput(
fantasy.ToolCallContent{
ToolCallID: "call-1",
ToolName: "exec",
Input: `{"command":":true"}`,
},
"Run the command 'echo dragonscale-eval-test' and tell me the output.",
)
if !strings.Contains(got.Input, `"echo dragonscale-eval-test"`) {
t.Fatalf("expected forced exec command, got %q", got.Input)
}
}
func TestRepairToolCallInputRepairsNestedToolCallPlaceholder(t *testing.T) {
t.Parallel()
got := repairToolCallInput(
fantasy.ToolCallContent{
ToolCallID: "call-1",
ToolName: "tool_call",
Input: `{"tool_name":"skill_read","arguments":{"name":":"}}`,
},
"Read the 'eval-test-skill' skill and tell me what greeting templates it provides.",
)
if !strings.Contains(got.Input, `"eval-test-skill"`) {
t.Fatalf("expected repaired nested skill name, got %q", got.Input)
}
}
func TestRepairToolCallInputRepairsWriteAndListPlaceholders(t *testing.T) {
t.Parallel()
write := repairToolCallInput(
fantasy.ToolCallContent{
ToolCallID: "call-write",
ToolName: "write_file",
Input: `{"path":"}","content":","}`,
},
"Create a file called project/readme.txt with 'Project initialized'. Then list the project directory to verify it exists.",
)
if !strings.Contains(write.Input, `"project/readme.txt"`) || !strings.Contains(write.Input, `"Project initialized"`) {
t.Fatalf("expected repaired write args, got %q", write.Input)
}
list := repairToolCallInput(
fantasy.ToolCallContent{
ToolCallID: "call-list",
ToolName: "list_dir",
Input: `{}`,
},
"Create a file called project/readme.txt with 'Project initialized'. Then list the project directory to verify it exists.",
)
if !strings.Contains(list.Input, `"project"`) {
t.Fatalf("expected repaired list_dir path, got %q", list.Input)
}
}
func TestRepairToolCallInputRepairsReadBackAndFetchPlaceholders(t *testing.T) {
t.Parallel()
read := repairToolCallInput(
fantasy.ToolCallContent{
ToolCallID: "call-read",
ToolName: "read_file",
Input: `{"path":"}"}`,
},
"Create a file called test_steps.txt with the content 'step test', then read it back to confirm.",
)
if !strings.Contains(read.Input, `"test_steps.txt"`) {
t.Fatalf("expected repaired read_file path, got %q", read.Input)
}
fetch := repairToolCallInput(
fantasy.ToolCallContent{
ToolCallID: "call-fetch",
ToolName: "web_fetch",
Input: `{"url":".example.com"}`,
},
"Fetch the contents of https://example.com and tell me the title of the page.",
)
if !strings.Contains(fetch.Input, `"https://example.com"`) {
t.Fatalf("expected repaired web_fetch url, got %q", fetch.Input)
}
}
func TestSanitizePolicyErrorPreservesSafeExecErrors(t *testing.T) {
t.Parallel()
raw := "command timed out after 8s"
if got := sanitizePolicyError(raw); got != raw {
t.Fatalf("expected timeout text to survive sanitization, got %q", got)
}
}
func TestSanitizePolicyErrorRedactsPolicyViolations(t *testing.T) {
t.Parallel()
got := sanitizePolicyError("filesystem access denied: /etc/passwd")
if got != "policy violation: filesystem access denied" {
t.Fatalf("expected redacted policy text, got %q", got)
}
}

View file

@ -0,0 +1,111 @@
package agent
import (
"os"
"path/filepath"
"testing"
jsonv2 "github.com/go-json-experiment/json"
pkgroot "github.com/ZanzyTHEbar/dragonscale/pkg"
"github.com/ZanzyTHEbar/dragonscale/pkg/bus"
"github.com/ZanzyTHEbar/dragonscale/pkg/config"
memstore "github.com/ZanzyTHEbar/dragonscale/pkg/memory/store"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestContextBuilder_UsesActiveSessionWorkingContext(t *testing.T) {
t.Parallel()
al := newSessionBoundAgentLoop(t)
require.NoError(t, al.memoryStore.SetWorkingContext(t.Context(), pkgroot.NAME, "session-a", "working context A"))
require.NoError(t, al.memoryStore.SetWorkingContext(t.Context(), pkgroot.NAME, "session-b", "working context B"))
al.activeSessionKey.Store("session-a")
promptA := al.contextBuilder.BuildSystemPromptWithBudget(0)
assert.Contains(t, promptA, "working context A")
assert.NotContains(t, promptA, "working context B")
al.activeSessionKey.Store("session-b")
promptB := al.contextBuilder.BuildSystemPromptWithBudget(0)
assert.Contains(t, promptB, "working context B")
assert.NotContains(t, promptB, "working context A")
}
func TestMemGPTTool_UsesResolvedSessionForWriteAndStatus(t *testing.T) {
t.Parallel()
al := newSessionBoundAgentLoop(t)
toolAny, ok := al.tools.Get("memory")
require.True(t, ok)
memTool, ok := toolAny.(*MemGPTTool)
require.True(t, ok)
al.activeSessionKey.Store("session-a")
writeA := memTool.Execute(t.Context(), map[string]interface{}{
"action": "write",
"content": "session-scoped memory",
"tier": "recall",
"sector": "semantic",
})
require.False(t, writeA.IsError)
al.activeSessionKey.Store("session-b")
writeB := memTool.Execute(t.Context(), map[string]interface{}{
"action": "write",
"content": "session-scoped memory",
"tier": "recall",
"sector": "semantic",
})
require.False(t, writeB.IsError)
sessionAItems, err := al.memDelegate.ListRecallItems(t.Context(), pkgroot.NAME, "session-a", 10, 0)
require.NoError(t, err)
sessionBItems, err := al.memDelegate.ListRecallItems(t.Context(), pkgroot.NAME, "session-b", 10, 0)
require.NoError(t, err)
require.NotEmpty(t, sessionAItems)
require.NotEmpty(t, sessionBItems)
assert.Equal(t, "session-scoped memory", sessionAItems[0].Content)
assert.Equal(t, "session-scoped memory", sessionBItems[0].Content)
require.NoError(t, al.memoryStore.SetWorkingContext(t.Context(), pkgroot.NAME, "session-a", "scoped working context"))
al.activeSessionKey.Store("session-a")
statusA := invokeMemoryAction(t, memTool, map[string]interface{}{"action": "status"})
require.NotNil(t, statusA.Status)
assert.Greater(t, statusA.Status.WorkingContextTokens, 0)
assert.Equal(t, 1, statusA.Status.RecallItemCount)
al.activeSessionKey.Store("session-b")
statusB := invokeMemoryAction(t, memTool, map[string]interface{}{"action": "status"})
require.NotNil(t, statusB.Status)
assert.Equal(t, 0, statusB.Status.WorkingContextTokens)
assert.Equal(t, 1, statusB.Status.RecallItemCount)
}
func newSessionBoundAgentLoop(t *testing.T) *AgentLoop {
t.Helper()
tmpDir, err := os.MkdirTemp("", "session-binding-*")
require.NoError(t, err)
t.Cleanup(func() { _ = os.RemoveAll(tmpDir) })
cfg := config.DefaultConfig()
cfg.Agents.Defaults.Sandbox = tmpDir
cfg.Memory.DBPath = filepath.Join(tmpDir, "session-binding.db")
return mustNewAgentLoop(t, cfg, bus.NewMessageBus(), newMockLanguageModel("ok"))
}
func invokeMemoryAction(t *testing.T, tool *MemGPTTool, args map[string]interface{}) memstore.MemoryToolResponse {
t.Helper()
result := tool.Execute(t.Context(), args)
require.False(t, result.IsError)
var response memstore.MemoryToolResponse
require.NoError(t, jsonv2.Unmarshal([]byte(result.ForLLM), &response))
require.True(t, response.Success)
return response
}

View file

@ -2,9 +2,9 @@ package agent
import ( import (
"context" "context"
"crypto/sha1"
"encoding/json" "encoding/json"
"fmt" "fmt"
"sort"
"strings" "strings"
"time" "time"
@ -102,6 +102,7 @@ func (al *AgentLoop) persistEmergencyProvenance(ctx context.Context, prov Emerge
Action: "emergency_compression", Action: "emergency_compression",
Target: fmt.Sprintf("cycle_%d", prov.Cycle), Target: fmt.Sprintf("cycle_%d", prov.Cycle),
Input: string(input), Input: string(input),
Success: true,
} }
aCtx, cancel := context.WithTimeout(ctx, time.Second) aCtx, cancel := context.WithTimeout(ctx, time.Second)
defer cancel() defer cancel()
@ -504,9 +505,12 @@ func (al *AgentLoop) sessionsToMessagePairs(sessionKey string) []observation.Mes
// contextTreeCacheEntry caches rendered query-selected history blocks per session. // contextTreeCacheEntry caches rendered query-selected history blocks per session.
type contextTreeCacheEntry struct { type contextTreeCacheEntry struct {
msgCount int msgCount int
query string query string
rendered string rendered string
accessCounts map[string]int
prevSelection map[string]float64
selectedKeys []string
} }
// applyContextTreeSelection selects relevant historical context via query-adaptive // applyContextTreeSelection selects relevant historical context via query-adaptive
@ -540,46 +544,205 @@ func (al *AgentLoop) applyContextTreeSelection(ctx context.Context, sessionKey,
} }
cacheHit := false cacheHit := false
priorEntry := contextTreeCacheEntry{}
if cached, ok := al.contextTreeCache.Load(sessionKey); ok { if cached, ok := al.contextTreeCache.Load(sessionKey); ok {
entry := cached.(contextTreeCacheEntry) priorEntry = cached.(contextTreeCacheEntry)
if entry.msgCount == len(compressible) && entry.query == query { if priorEntry.msgCount == len(compressible) && priorEntry.query == query {
al.contextBuilder.SetContextTreeBlock(entry.rendered) cacheHit = true
al.contextBuilder.SetContextTreeBlock(priorEntry.rendered)
if len(priorEntry.selectedKeys) > 0 {
updated := contextTreeCacheEntry{
msgCount: priorEntry.msgCount,
query: priorEntry.query,
rendered: priorEntry.rendered,
accessCounts: cloneContextTreeAccessCounts(priorEntry.accessCounts),
prevSelection: cloneContextTreeSelection(priorEntry.prevSelection),
selectedKeys: append([]string(nil), priorEntry.selectedKeys...),
}
for _, key := range updated.selectedKeys {
updated.accessCounts[key]++
}
al.contextTreeCache.Store(sessionKey, updated)
}
return tail return tail
} }
} }
tree := contexttree.NewContextTree(contexttree.DefaultScoringConfig()) tree := contexttree.NewContextTree(contexttree.DefaultScoringConfig())
rootID := tree.Root.ID rootID := tree.Root.ID
for _, m := range compressible {
tree.AddNode(rootID, contextNodeTypeForRole(m.Role), m.Content, nil, contexttree.ExtractTerms(m.Content))
}
queryTerms := contexttree.ExtractTerms(query) queryText := strings.TrimSpace(query)
if queryText == "" && len(tail) > 0 {
queryText = strings.TrimSpace(tail[len(tail)-1].Content)
}
queryTerms := contexttree.ExtractTerms(queryText)
if len(queryTerms) == 0 && len(tail) > 0 { if len(queryTerms) == 0 && len(tail) > 0 {
queryTerms = contexttree.ExtractTerms(tail[len(tail)-1].Content) queryTerms = contexttree.ExtractTerms(tail[len(tail)-1].Content)
} }
scores := tree.ScoreAll(nil, queryTerms) var (
nodes := make([]*contexttree.ContextNode, 0, len(tree.NodeIndex)-1) queryEmbedding []float32
for id, node := range tree.NodeIndex { nodeEmbeddings []memory.Embedding
if node.Type == contexttree.NodeTypeRoot { )
continue if al.memoryStore != nil && al.memoryStore.Embedder() != nil {
embedder := al.memoryStore.Embedder()
if queryText != "" {
if embedded, err := embedder.Embed(ctx, queryText); err == nil {
queryEmbedding = embedded
} else {
logger.WarnCF("agent", "Context-Tree query embedding failed",
map[string]interface{}{"error": err.Error(), "session": sessionKey})
}
}
texts := make([]string, 0, len(compressible))
for _, m := range compressible {
texts = append(texts, m.Content)
}
if len(texts) > 0 {
if embedded, err := embedder.EmbedBatch(ctx, texts); err == nil {
nodeEmbeddings = embedded
} else {
logger.WarnCF("agent", "Context-Tree batch embedding failed",
map[string]interface{}{"error": err.Error(), "session": sessionKey, "messages": len(texts)})
}
} }
node.TotalScore = scores[id]
nodes = append(nodes, node)
} }
sort.Slice(nodes, func(i, j int) bool { stableKeys := make(map[string]*contexttree.ContextNode, len(compressible))
if nodes[i].TotalScore == nodes[j].TotalScore { nodeStableKeys := make(map[ids.UUID]string, len(compressible))
return nodes[i].CreatedAt.After(nodes[j].CreatedAt) for idx, m := range compressible {
var embedding []float32
if idx < len(nodeEmbeddings) {
embedding = nodeEmbeddings[idx]
} }
return nodes[i].TotalScore > nodes[j].TotalScore stableKey := contextTreeStableKey(idx, m)
}) node := tree.AddNode(rootID, contextNodeTypeForRole(m.Role), m.Content, embedding, contexttree.ExtractTerms(m.Content))
if priorCount := priorEntry.accessCounts[stableKey]; priorCount > 1 {
node.AccessCount = priorCount
}
stableKeys[stableKey] = node
nodeStableKeys[node.ID] = stableKey
}
selectionBudget := budget.DAGSummaries selectionBudget := budget.DAGSummaries
if selectionBudget <= 0 { if selectionBudget <= 0 {
selectionBudget = 512 selectionBudget = 512
} }
prevSelection := make(map[ids.UUID]float64, len(priorEntry.prevSelection))
for stableKey, score := range priorEntry.prevSelection {
if node, ok := stableKeys[stableKey]; ok {
prevSelection[node.ID] = score
}
}
budgetCount := contextTreeSelectionBudgetCount(compressible, selectionBudget)
stickySelected := tree.SelectNodesWithHysteresis(queryEmbedding, queryTerms, budgetCount, prevSelection)
sampledSelected := tree.PruneWithTemperature(queryEmbedding, queryTerms, budgetCount)
selected := mergeContextTreeSelections(stickySelected, sampledSelected)
selected = fitContextTreeSelectionToTokenBudget(selected, selectionBudget)
usedTokens := 0
for _, node := range selected {
tree.RecordAccess(node.ID)
usedTokens += observation.EstimateTokens(node.Content)
}
rendered := renderContextTreeSelection(selected)
al.contextBuilder.SetContextTreeBlock(rendered)
nextEntry := contextTreeCacheEntry{
msgCount: len(compressible),
query: query,
rendered: rendered,
accessCounts: cloneContextTreeAccessCounts(priorEntry.accessCounts),
prevSelection: make(map[string]float64, len(selected)),
selectedKeys: make([]string, 0, len(selected)),
}
for _, node := range selected {
stableKey, ok := nodeStableKeys[node.ID]
if !ok {
continue
}
nextEntry.selectedKeys = append(nextEntry.selectedKeys, stableKey)
nextEntry.prevSelection[stableKey] = node.TotalScore
nextEntry.accessCounts[stableKey] = node.AccessCount
}
al.contextTreeCache.Store(sessionKey, nextEntry)
logger.DebugCF("agent", "Context-Tree selection applied", map[string]interface{}{
"total_msgs": len(history),
"compressed_msgs": len(compressible),
"tail_msgs": len(tail),
"selected_nodes": len(selected),
"selected_tokens": usedTokens,
"cache_hit": cacheHit,
"semantic_enabled": len(queryEmbedding) > 0 && len(nodeEmbeddings) == len(compressible),
"sticky_candidates": len(stickySelected),
"sampled_nodes": len(sampledSelected),
})
return tail
}
func contextTreeStableKey(index int, msg messages.Message) string {
sum := sha1.Sum([]byte(msg.Role + "\x00" + msg.Content))
return fmt.Sprintf("%06d:%s:%x", index, msg.Role, sum[:6])
}
func contextTreeSelectionBudgetCount(history []messages.Message, tokenBudget int) int {
if len(history) == 0 {
return 0
}
if tokenBudget <= 0 {
return len(history)
}
totalTokens := 0
for _, msg := range history {
totalTokens += observation.EstimateTokens(msg.Content)
}
avgTokens := 64
if totalTokens > 0 {
avgTokens = max(32, totalTokens/len(history))
}
count := tokenBudget / avgTokens
if count < 1 {
count = 1
}
if count > len(history) {
count = len(history)
}
return count
}
func mergeContextTreeSelections(sticky, sampled []*contexttree.ContextNode) []*contexttree.ContextNode {
seen := make(map[ids.UUID]struct{}, len(sticky)+len(sampled))
merged := make([]*contexttree.ContextNode, 0, len(sticky)+len(sampled))
for _, group := range [][]*contexttree.ContextNode{sticky, sampled} {
for _, node := range group {
if node == nil {
continue
}
if _, ok := seen[node.ID]; ok {
continue
}
seen[node.ID] = struct{}{}
merged = append(merged, node)
}
}
return merged
}
func fitContextTreeSelectionToTokenBudget(nodes []*contexttree.ContextNode, budget int) []*contexttree.ContextNode {
if budget <= 0 || len(nodes) == 0 {
return nil
}
selected := make([]*contexttree.ContextNode, 0, len(nodes)) selected := make([]*contexttree.ContextNode, 0, len(nodes))
usedTokens := 0 usedTokens := 0
for _, node := range nodes { for _, node := range nodes {
@ -587,27 +750,40 @@ func (al *AgentLoop) applyContextTreeSelection(ctx context.Context, sessionKey,
if nodeTokens == 0 { if nodeTokens == 0 {
continue continue
} }
if usedTokens+nodeTokens > selectionBudget { if usedTokens+nodeTokens > budget {
continue continue
} }
selected = append(selected, node) selected = append(selected, node)
usedTokens += nodeTokens usedTokens += nodeTokens
} }
rendered := renderContextTreeSelection(selected) if len(selected) == 0 && len(nodes) > 0 {
al.contextBuilder.SetContextTreeBlock(rendered) return []*contexttree.ContextNode{nodes[0]}
al.contextTreeCache.Store(sessionKey, contextTreeCacheEntry{msgCount: len(compressible), query: query, rendered: rendered}) }
logger.DebugCF("agent", "Context-Tree selection applied", map[string]interface{}{ return selected
"total_msgs": len(history), }
"compressed_msgs": len(compressible),
"tail_msgs": len(tail),
"selected_nodes": len(selected),
"selected_tokens": usedTokens,
"cache_hit": cacheHit,
})
return tail func cloneContextTreeAccessCounts(src map[string]int) map[string]int {
if len(src) == 0 {
return make(map[string]int)
}
dst := make(map[string]int, len(src))
for key, value := range src {
dst[key] = value
}
return dst
}
func cloneContextTreeSelection(src map[string]float64) map[string]float64 {
if len(src) == 0 {
return make(map[string]float64)
}
dst := make(map[string]float64, len(src))
for key, value := range src {
dst[key] = value
}
return dst
} }
func contextNodeTypeForRole(role string) contexttree.NodeType { func contextNodeTypeForRole(role string) contexttree.NodeType {

View file

@ -2,38 +2,20 @@ package agent
import ( import (
"context" "context"
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg" "github.com/ZanzyTHEbar/dragonscale/pkg"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids" "github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/logger" "github.com/ZanzyTHEbar/dragonscale/pkg/logger"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
) )
// TaskCompletion tracks the outcome of an agent run for RL analysis. type TaskCompletion = memory.TaskCompletionRecord
type TaskCompletion struct { type MemoryRating = memory.MemoryRating
TaskID string
Description string
TokensUsed int
ToolCalls int
Errors int
UserCorrections int
Completed bool
SelfReports []MemoryRating // MemoryID + Score
CreatedAt time.Time
}
// MemoryRating represents a self-reported usefulness score for a memory.
type MemoryRating struct {
MemoryID ids.UUID
Score int // 0-3 scale
}
// TaskCompletionStore is the interface for storing task completion records. // TaskCompletionStore is the interface for storing task completion records.
// Implemented by the memory delegate. // Implemented by the memory delegate.
type TaskCompletionStore interface { type TaskCompletionStore interface {
StoreTaskCompletion(ctx context.Context, agentID string, completion TaskCompletion, conversationID, runID ids.UUID) error StoreTaskCompletion(ctx context.Context, agentID string, completion TaskCompletion, conversationID, runID ids.UUID) error
GetCompletedTasks(ctx context.Context, agentID string, since time.Time) ([]TaskCompletion, error)
UpdateMemorySelfReport(ctx context.Context, memoryID ids.UUID, score int) error
} }
// endTask stores task completion data and self-reports. // endTask stores task completion data and self-reports.

View file

@ -58,6 +58,7 @@ func MakeUnifiedRunLoopFunc(al *AgentLoop) tools.RunLoopFunc {
Base: baseRuntime, Base: baseRuntime,
Bus: al.secureBus, Bus: al.secureBus,
SessionKey: sessionKey, SessionKey: sessionKey,
UserPrompt: userPrompt,
StateStore: al.stateStore, StateStore: al.stateStore,
RunID: runID, RunID: runID,
} }

View file

@ -30,14 +30,15 @@ type AuditAnalysisStore interface {
// AuditEntry represents a single audit log entry for analysis. // AuditEntry represents a single audit log entry for analysis.
type AuditEntry struct { type AuditEntry struct {
ID string ID string
Timestamp time.Time Timestamp time.Time
ToolName string ToolName string
ToolInput string ToolCallID string
Success bool ToolInput string
ErrorMsg string Success bool
SessionID string ErrorMsg string
AgentID string SessionID string
AgentID string
} }
// ToolSequence represents a tool call in a session sequence. // ToolSequence represents a tool call in a session sequence.

View file

@ -98,7 +98,12 @@ func DisableFileLogging() {
} }
func logMessage(level LogLevel, component string, message string, fields map[string]interface{}) { func logMessage(level LogLevel, component string, message string, fields map[string]interface{}) {
if level < currentLevel { mu.RLock()
current := currentLevel
file := logger.file
mu.RUnlock()
if level < current {
return return
} }
@ -117,10 +122,10 @@ func logMessage(level LogLevel, component string, message string, fields map[str
} }
} }
if logger.file != nil { if file != nil {
jsonData, err := jsonv2.Marshal(entry) jsonData, err := jsonv2.Marshal(entry)
if err == nil { if err == nil {
logger.file.WriteString(string(jsonData) + "\n") file.WriteString(string(jsonData) + "\n")
} }
} }

View file

@ -45,14 +45,15 @@ type RetrievedMemoryRecord struct {
// AuditEntry represents a single audit log entry for analysis. // AuditEntry represents a single audit log entry for analysis.
// Mirrors cortex.AuditEntry. // Mirrors cortex.AuditEntry.
type AuditEntry struct { type AuditEntry struct {
ID string ID string
Timestamp time.Time Timestamp time.Time
ToolName string ToolName string
ToolInput string ToolCallID string
Success bool ToolInput string
ErrorMsg string Success bool
SessionID string ErrorMsg string
AgentID string SessionID string
AgentID string
} }
// DetectedPattern represents a pattern detected from audit analysis. // DetectedPattern represents a pattern detected from audit analysis.

View file

@ -821,8 +821,11 @@ func (d *LibSQLDelegate) InsertAuditEntry(ctx context.Context, entry *memory.Aud
SessionKey: entry.SessionKey, SessionKey: entry.SessionKey,
Action: entry.Action, Action: entry.Action,
Target: entry.Target, Target: entry.Target,
ToolCallID: entry.ToolCallID,
Input: &entry.Input, Input: &entry.Input,
Output: &entry.Output, Output: &entry.Output,
Success: entry.Success,
ErrorMsg: entry.ErrorMsg,
DurationMs: ptrInt64(int64(entry.DurationMS)), DurationMs: ptrInt64(int64(entry.DurationMS)),
}) })
if err != nil { if err != nil {
@ -854,8 +857,11 @@ func (d *LibSQLDelegate) InsertAuditEntryBatch(ctx context.Context, entries []*m
SessionKey: entry.SessionKey, SessionKey: entry.SessionKey,
Action: entry.Action, Action: entry.Action,
Target: entry.Target, Target: entry.Target,
ToolCallID: entry.ToolCallID,
Input: &entry.Input, Input: &entry.Input,
Output: &entry.Output, Output: &entry.Output,
Success: entry.Success,
ErrorMsg: entry.ErrorMsg,
DurationMs: ptrInt64(int64(entry.DurationMS)), DurationMs: ptrInt64(int64(entry.DurationMS)),
}) })
if err != nil { if err != nil {
@ -1097,6 +1103,31 @@ func (d *LibSQLDelegate) UpdateMemoryWeight(ctx context.Context, memoryID ids.UU
}) })
} }
// StoreTaskCompletion persists a completed run for downstream RL analysis.
func (d *LibSQLDelegate) StoreTaskCompletion(ctx context.Context, agentID string, completion memory.TaskCompletionRecord, conversationID, runID ids.UUID) error {
tokensUsed := int64(completion.TokensUsed)
toolCalls := int64(completion.ToolCalls)
errorsCount := int64(completion.Errors)
userCorrections := int64(completion.UserCorrections)
_, err := d.queries.StoreTaskCompletion(ctx, memsqlc.StoreTaskCompletionParams{
ID: ids.New(),
AgentID: agentID,
ConversationID: conversationID,
RunID: runID,
Description: completion.Description,
TokensUsed: &tokensUsed,
ToolCalls: &toolCalls,
Errors: &errorsCount,
UserCorrections: &userCorrections,
Completed: completion.Completed,
})
if err != nil {
return fmt.Errorf("store task completion: %w", err)
}
return nil
}
// UpdateMemorySelfReport updates the self-reported score for a memory. // UpdateMemorySelfReport updates the self-reported score for a memory.
// Implements cortex.RLStore interface. // Implements cortex.RLStore interface.
func (d *LibSQLDelegate) UpdateMemorySelfReport(ctx context.Context, memoryID ids.UUID, score int) error { func (d *LibSQLDelegate) UpdateMemorySelfReport(ctx context.Context, memoryID ids.UUID, score int) error {
@ -1235,33 +1266,32 @@ func (d *LibSQLDelegate) GetRecentAuditEntries(ctx context.Context, since time.T
for _, row := range rows { for _, row := range rows {
lowerAction := strings.ToLower(strings.TrimSpace(row.Action)) lowerAction := strings.ToLower(strings.TrimSpace(row.Action))
toolName := strings.TrimSpace(row.Action) toolName := strings.TrimSpace(row.Action)
if strings.HasPrefix(lowerAction, "tool_") && strings.TrimSpace(row.Target) != "" { switch {
case strings.HasPrefix(lowerAction, "tool_"),
strings.HasPrefix(lowerAction, "memory_"),
strings.HasPrefix(lowerAction, "doc_"),
strings.HasPrefix(lowerAction, "state_"),
lowerAction == "emergency_compression":
toolName = strings.TrimSpace(row.Target) toolName = strings.TrimSpace(row.Target)
} }
if toolName == "" { if toolName == "" {
toolName = strings.TrimSpace(row.Target) toolName = strings.TrimSpace(row.Target)
} }
success := true
if lowerAction == "tool_error" || strings.Contains(lowerAction, "error") || strings.Contains(lowerAction, "fail") {
success = false
}
entry := AuditEntry{ entry := AuditEntry{
ID: row.ID.String(), ID: row.ID.String(),
Timestamp: row.CreatedAt, Timestamp: row.CreatedAt,
ToolName: toolName, ToolName: toolName,
ToolInput: "", ToolCallID: row.ToolCallID,
Success: success, ToolInput: "",
SessionID: row.SessionKey, Success: row.Success,
AgentID: row.AgentID, ErrorMsg: row.ErrorMsg,
SessionID: row.SessionKey,
AgentID: row.AgentID,
} }
if row.Input != nil { if row.Input != nil {
entry.ToolInput = *row.Input entry.ToolInput = *row.Input
} }
if !success && row.Output != nil {
entry.ErrorMsg = *row.Output
}
entries = append(entries, entry) entries = append(entries, entry)
} }
@ -1502,6 +1532,9 @@ func sqlcAuditToMemory(row memsqlc.AgentAuditLog) *memory.AuditEntry {
SessionKey: row.SessionKey, SessionKey: row.SessionKey,
Action: row.Action, Action: row.Action,
Target: row.Target, Target: row.Target,
ToolCallID: row.ToolCallID,
Success: row.Success,
ErrorMsg: row.ErrorMsg,
CreatedAt: row.CreatedAt, CreatedAt: row.CreatedAt,
} }
if row.Input != nil { if row.Input != nil {

View file

@ -2,6 +2,7 @@ package delegate
import ( import (
"context" "context"
"strings"
"testing" "testing"
"time" "time"
@ -13,14 +14,20 @@ import (
) )
func makeAuditEntry(agentID, sessionKey, action, target string) *memory.AuditEntry { func makeAuditEntry(agentID, sessionKey, action, target string) *memory.AuditEntry {
toolCallID := ""
if strings.HasPrefix(action, "tool") {
toolCallID = "call-" + target
}
return &memory.AuditEntry{ return &memory.AuditEntry{
ID: ids.New(), ID: ids.New(),
AgentID: agentID, AgentID: agentID,
SessionKey: sessionKey, SessionKey: sessionKey,
Action: action, Action: action,
Target: target, Target: target,
ToolCallID: toolCallID,
Input: `{"arg":"val"}`, Input: `{"arg":"val"}`,
Output: `{"result":"ok"}`, Output: `{"result":"ok"}`,
Success: true,
DurationMS: 42, DurationMS: 42,
} }
} }
@ -66,6 +73,86 @@ func TestLibSQLDelegate_InsertAuditEntry(t *testing.T) {
} }
} }
func TestLibSQLDelegate_ListAuditEntries_PreservesOutcomeMetadata(t *testing.T) {
t.Parallel()
d := newTestDelegate(t)
ctx := t.Context()
entry := &memory.AuditEntry{
ID: ids.New(),
AgentID: "a1",
SessionKey: "sess-1",
Action: "tool_error",
Target: "exec",
ToolCallID: "call-exec",
Input: `{"command":"rm -rf /tmp/test"}`,
Output: "permission denied",
Success: false,
ErrorMsg: "permission denied",
DurationMS: 9,
}
require.NoError(t, d.InsertAuditEntry(ctx, entry))
entries, err := d.ListAuditEntries(ctx, "a1", 10)
require.NoError(t, err)
require.Len(t, entries, 1)
assert.Equal(t, "call-exec", entries[0].ToolCallID)
assert.False(t, entries[0].Success)
assert.Equal(t, "permission denied", entries[0].ErrorMsg)
}
func TestLibSQLDelegate_GetRecentAuditEntries_HandlesMixedLegacyAndExplicitRows(t *testing.T) {
t.Parallel()
d := newTestDelegate(t)
ctx := t.Context()
now := time.Now().UTC()
_, err := d.db.ExecContext(ctx, `
INSERT INTO agent_audit_log (
id, agent_id, session_key, action, target, input, output, duration_ms, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
`, ids.New(), "a1", "sess-legacy", "tool_call", "legacy_read", `{"path":"legacy.txt"}`, `{"status":"ok"}`, 4, now)
require.NoError(t, err)
explicit := &memory.AuditEntry{
ID: ids.New(),
AgentID: "a1",
SessionKey: "sess-explicit",
Action: "tool_call",
Target: "explicit_write",
ToolCallID: "call-explicit-write",
Input: `{"path":"new.txt"}`,
Output: "write denied",
Success: false,
ErrorMsg: "write denied",
DurationMS: 7,
}
require.NoError(t, d.InsertAuditEntry(ctx, explicit))
entries, err := d.GetRecentAuditEntries(ctx, now.Add(-time.Minute))
require.NoError(t, err)
require.Len(t, entries, 2)
byTool := make(map[string]AuditEntry, len(entries))
for _, entry := range entries {
byTool[entry.ToolName] = entry
}
legacy, ok := byTool["legacy_read"]
require.True(t, ok, "legacy row missing")
assert.True(t, legacy.Success)
assert.Equal(t, "", legacy.ErrorMsg)
explicitEntry, ok := byTool["explicit_write"]
require.True(t, ok, "explicit row missing")
assert.False(t, explicitEntry.Success)
assert.Equal(t, "write denied", explicitEntry.ErrorMsg)
assert.Equal(t, "call-explicit-write", explicitEntry.ToolCallID)
assert.Equal(t, `{"path":"new.txt"}`, explicitEntry.ToolInput)
}
func TestLibSQLDelegate_ListAuditEntries(t *testing.T) { func TestLibSQLDelegate_ListAuditEntries(t *testing.T) {
t.Parallel() t.Parallel()
tests := []struct { tests := []struct {

View file

@ -84,7 +84,9 @@ func BenchmarkInsertAuditEntry(b *testing.B) {
SessionKey: "bench-sess", SessionKey: "bench-sess",
Action: "tool_call", Action: "tool_call",
Target: "read_file", Target: "read_file",
ToolCallID: "call-read-file",
Input: `{"path": "/tmp/test"}`, Input: `{"path": "/tmp/test"}`,
Success: true,
}) })
} }
} }

View file

@ -156,8 +156,10 @@ func TestEndToEnd_SessionAndAuditFlow(t *testing.T) {
SessionKey: sessionKey, SessionKey: sessionKey,
Action: "tool_call", Action: "tool_call",
Target: "weather_api", Target: "weather_api",
ToolCallID: "call-weather-api",
Input: `{"location":"here"}`, Input: `{"location":"here"}`,
Output: `{"temp":72}`, Output: `{"temp":72}`,
Success: true,
DurationMS: 150, DurationMS: 150,
} }
require.NoError(t, d.InsertAuditEntry(ctx, auditEntry)) require.NoError(t, d.InsertAuditEntry(ctx, auditEntry))

View file

@ -690,7 +690,9 @@ func TestIntegration_FullStackNoDisk(t *testing.T) {
SessionKey: "test-session", SessionKey: "test-session",
Action: "tool_call", Action: "tool_call",
Target: "exec", Target: "exec",
ToolCallID: "call-exec",
Input: `{"command":"ls"}`, Input: `{"command":"ls"}`,
Success: true,
} }
if err := d.InsertAuditEntry(ctx, entry); err != nil { if err := d.InsertAuditEntry(ctx, entry); err != nil {
t.Fatalf("InsertAuditEntry: %v", err) t.Fatalf("InsertAuditEntry: %v", err)

View file

@ -293,8 +293,11 @@ type AuditEntry struct {
SessionKey string SessionKey string
Action string // "tool_call", "memory_write", "doc_update", "state_change" Action string // "tool_call", "memory_write", "doc_update", "state_change"
Target string // tool name, doc name, key name Target string // tool name, doc name, key name
ToolCallID string
Input string Input string
Output string Output string
Success bool
ErrorMsg string
DurationMS int DurationMS int
CreatedAt time.Time CreatedAt time.Time
} }

View file

@ -0,0 +1,48 @@
package migrations
import (
"context"
"database/sql"
"fmt"
"github.com/pressly/goose/v3"
)
func init() {
goose.AddMigrationContext(up017AgentAuditOutcomes, down017AgentAuditOutcomes)
}
func up017AgentAuditOutcomes(ctx context.Context, tx *sql.Tx) error {
stmts := []string{
`ALTER TABLE agent_audit_log ADD COLUMN success BOOLEAN NOT NULL DEFAULT TRUE`,
`ALTER TABLE agent_audit_log ADD COLUMN error_msg TEXT NOT NULL DEFAULT ''`,
`ALTER TABLE agent_audit_log ADD COLUMN tool_call_id TEXT NOT NULL DEFAULT ''`,
`UPDATE agent_audit_log
SET success = CASE
WHEN lower(action) = 'tool_error'
OR instr(lower(action), 'error') > 0
OR instr(lower(action), 'fail') > 0
THEN FALSE
ELSE TRUE
END`,
`UPDATE agent_audit_log
SET error_msg = COALESCE(output, '')
WHERE success = FALSE
AND error_msg = ''`,
`CREATE INDEX IF NOT EXISTS idx_audit_tool_call_id ON agent_audit_log(tool_call_id)`,
}
for _, s := range stmts {
if _, err := tx.ExecContext(ctx, s); err != nil {
return fmt.Errorf("017_agent_audit_outcomes up: %w\nSQL: %s", err, s)
}
}
return nil
}
func down017AgentAuditOutcomes(ctx context.Context, tx *sql.Tx) error {
if _, err := tx.ExecContext(ctx, `DROP INDEX IF EXISTS idx_audit_tool_call_id`); err != nil {
return fmt.Errorf("017_agent_audit_outcomes down: %w", err)
}
return nil
}

View file

@ -66,8 +66,11 @@ INSERT INTO agent_audit_log (
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
) )
@ -80,6 +83,9 @@ VALUES (
?6, ?6,
?7, ?7,
?8, ?8,
?9,
?10,
?11,
strftime('%Y-%m-%dT%H:%M:%fZ', 'now') strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
) )
RETURNING id, RETURNING id,
@ -87,8 +93,11 @@ RETURNING id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
` `
@ -99,8 +108,11 @@ type InsertAuditEntryParams struct {
SessionKey string `db:"session_key" json:"session_key"` SessionKey string `db:"session_key" json:"session_key"`
Action string `db:"action" json:"action"` Action string `db:"action" json:"action"`
Target string `db:"target" json:"target"` Target string `db:"target" json:"target"`
ToolCallID string `db:"tool_call_id" json:"tool_call_id"`
Input *string `db:"input" json:"input"` Input *string `db:"input" json:"input"`
Output *string `db:"output" json:"output"` Output *string `db:"output" json:"output"`
Success bool `db:"success" json:"success"`
ErrorMsg string `db:"error_msg" json:"error_msg"`
DurationMs *int64 `db:"duration_ms" json:"duration_ms"` DurationMs *int64 `db:"duration_ms" json:"duration_ms"`
} }
@ -112,8 +124,11 @@ type InsertAuditEntryParams struct {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// ) // )
@ -126,6 +141,9 @@ type InsertAuditEntryParams struct {
// ?6, // ?6,
// ?7, // ?7,
// ?8, // ?8,
// ?9,
// ?10,
// ?11,
// strftime('%Y-%m-%dT%H:%M:%fZ', 'now') // strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
// ) // )
// RETURNING id, // RETURNING id,
@ -133,8 +151,11 @@ type InsertAuditEntryParams struct {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
func (q *Queries) InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) (AgentAuditLog, error) { func (q *Queries) InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) (AgentAuditLog, error) {
@ -144,8 +165,11 @@ func (q *Queries) InsertAuditEntry(ctx context.Context, arg InsertAuditEntryPara
arg.SessionKey, arg.SessionKey,
arg.Action, arg.Action,
arg.Target, arg.Target,
arg.ToolCallID,
arg.Input, arg.Input,
arg.Output, arg.Output,
arg.Success,
arg.ErrorMsg,
arg.DurationMs, arg.DurationMs,
) )
var i AgentAuditLog var i AgentAuditLog
@ -155,8 +179,11 @@ func (q *Queries) InsertAuditEntry(ctx context.Context, arg InsertAuditEntryPara
&i.SessionKey, &i.SessionKey,
&i.Action, &i.Action,
&i.Target, &i.Target,
&i.ToolCallID,
&i.Input, &i.Input,
&i.Output, &i.Output,
&i.Success,
&i.ErrorMsg,
&i.DurationMs, &i.DurationMs,
&i.CreatedAt, &i.CreatedAt,
) )
@ -169,8 +196,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -191,8 +221,11 @@ type ListAuditEntriesParams struct {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -214,8 +247,11 @@ func (q *Queries) ListAuditEntries(ctx context.Context, arg ListAuditEntriesPara
&i.SessionKey, &i.SessionKey,
&i.Action, &i.Action,
&i.Target, &i.Target,
&i.ToolCallID,
&i.Input, &i.Input,
&i.Output, &i.Output,
&i.Success,
&i.ErrorMsg,
&i.DurationMs, &i.DurationMs,
&i.CreatedAt, &i.CreatedAt,
); err != nil { ); err != nil {
@ -238,8 +274,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -262,8 +301,11 @@ type ListAuditEntriesByActionParams struct {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -286,8 +328,11 @@ func (q *Queries) ListAuditEntriesByAction(ctx context.Context, arg ListAuditEnt
&i.SessionKey, &i.SessionKey,
&i.Action, &i.Action,
&i.Target, &i.Target,
&i.ToolCallID,
&i.Input, &i.Input,
&i.Output, &i.Output,
&i.Success,
&i.ErrorMsg,
&i.DurationMs, &i.DurationMs,
&i.CreatedAt, &i.CreatedAt,
); err != nil { ); err != nil {
@ -310,8 +355,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -334,8 +382,11 @@ type ListAuditEntriesBySessionParams struct {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -358,8 +409,11 @@ func (q *Queries) ListAuditEntriesBySession(ctx context.Context, arg ListAuditEn
&i.SessionKey, &i.SessionKey,
&i.Action, &i.Action,
&i.Target, &i.Target,
&i.ToolCallID,
&i.Input, &i.Input,
&i.Output, &i.Output,
&i.Success,
&i.ErrorMsg,
&i.DurationMs, &i.DurationMs,
&i.CreatedAt, &i.CreatedAt,
); err != nil { ); err != nil {
@ -382,8 +436,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -402,8 +459,11 @@ type ListAuditEntriesGlobalParams struct {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -424,8 +484,11 @@ func (q *Queries) ListAuditEntriesGlobal(ctx context.Context, arg ListAuditEntri
&i.SessionKey, &i.SessionKey,
&i.Action, &i.Action,
&i.Target, &i.Target,
&i.ToolCallID,
&i.Input, &i.Input,
&i.Output, &i.Output,
&i.Success,
&i.ErrorMsg,
&i.DurationMs, &i.DurationMs,
&i.CreatedAt, &i.CreatedAt,
); err != nil { ); err != nil {
@ -448,13 +511,17 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
WHERE julianday(created_at) > julianday(?1) WHERE julianday(created_at) > julianday(?1)
ORDER BY created_at ASC, id ASC ORDER BY created_at ASC,
id ASC
LIMIT ?3 OFFSET ?2 LIMIT ?3 OFFSET ?2
` `
@ -471,13 +538,17 @@ type ListAuditEntriesGlobalSincePagedParams struct {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
// WHERE julianday(created_at) > julianday(?1) // WHERE julianday(created_at) > julianday(?1)
// ORDER BY created_at ASC, id ASC // ORDER BY created_at ASC,
// id ASC
// LIMIT ?3 OFFSET ?2 // LIMIT ?3 OFFSET ?2
func (q *Queries) ListAuditEntriesGlobalSincePaged(ctx context.Context, arg ListAuditEntriesGlobalSincePagedParams) ([]AgentAuditLog, error) { func (q *Queries) ListAuditEntriesGlobalSincePaged(ctx context.Context, arg ListAuditEntriesGlobalSincePagedParams) ([]AgentAuditLog, error) {
rows, err := q.db.QueryContext(ctx, ListAuditEntriesGlobalSincePaged, arg.Since, arg.Off, arg.Lim) rows, err := q.db.QueryContext(ctx, ListAuditEntriesGlobalSincePaged, arg.Since, arg.Off, arg.Lim)
@ -494,8 +565,11 @@ func (q *Queries) ListAuditEntriesGlobalSincePaged(ctx context.Context, arg List
&i.SessionKey, &i.SessionKey,
&i.Action, &i.Action,
&i.Target, &i.Target,
&i.ToolCallID,
&i.Input, &i.Input,
&i.Output, &i.Output,
&i.Success,
&i.ErrorMsg,
&i.DurationMs, &i.DurationMs,
&i.CreatedAt, &i.CreatedAt,
); err != nil { ); err != nil {

View file

@ -68,6 +68,37 @@ func (q *Queries) GetAgentConversation(ctx context.Context, arg GetAgentConversa
return i, err return i, err
} }
const GetLatestAgentConversationByTitle = `-- name: GetLatestAgentConversationByTitle :one
SELECT id, title, created_at, updated_at
FROM agent_conversations
WHERE title = ?
ORDER BY created_at DESC
LIMIT 1
`
type GetLatestAgentConversationByTitleParams struct {
Title *string `db:"title" json:"title"`
}
// GetLatestAgentConversationByTitle
//
// SELECT id, title, created_at, updated_at
// FROM agent_conversations
// WHERE title = ?
// ORDER BY created_at DESC
// LIMIT 1
func (q *Queries) GetLatestAgentConversationByTitle(ctx context.Context, arg GetLatestAgentConversationByTitleParams) (AgentConversation, error) {
row := q.db.QueryRowContext(ctx, GetLatestAgentConversationByTitle, arg.Title)
var i AgentConversation
err := row.Scan(
&i.ID,
&i.Title,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const ListAgentConversations = `-- name: ListAgentConversations :many const ListAgentConversations = `-- name: ListAgentConversations :many
SELECT id, title, created_at, updated_at SELECT id, title, created_at, updated_at
FROM agent_conversations FROM agent_conversations

View file

@ -18,8 +18,11 @@ type AgentAuditLog struct {
SessionKey string `db:"session_key" json:"session_key"` SessionKey string `db:"session_key" json:"session_key"`
Action string `db:"action" json:"action"` Action string `db:"action" json:"action"`
Target string `db:"target" json:"target"` Target string `db:"target" json:"target"`
ToolCallID string `db:"tool_call_id" json:"tool_call_id"`
Input *string `db:"input" json:"input"` Input *string `db:"input" json:"input"`
Output *string `db:"output" json:"output"` Output *string `db:"output" json:"output"`
Success bool `db:"success" json:"success"`
ErrorMsg string `db:"error_msg" json:"error_msg"`
DurationMs *int64 `db:"duration_ms" json:"duration_ms"` DurationMs *int64 `db:"duration_ms" json:"duration_ms"`
CreatedAt time.Time `db:"created_at" json:"created_at"` CreatedAt time.Time `db:"created_at" json:"created_at"`
} }

View file

@ -187,6 +187,7 @@ type Querier interface {
// WHERE agent_id = ?1 // WHERE agent_id = ?1
// AND session_key = ?2 // AND session_key = ?2
// AND tags = 'session-message' // AND tags = 'session-message'
// AND suppressed_at IS NULL
CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error) CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error)
//CreateAgentCheckpoint //CreateAgentCheckpoint
// //
@ -506,14 +507,14 @@ type Querier interface {
GetDocument(ctx context.Context, arg GetDocumentParams) (AgentDocument, error) GetDocument(ctx context.Context, arg GetDocumentParams) (AgentDocument, error)
// Get sessions with high token usage grouped by conversation/agent // Get sessions with high token usage grouped by conversation/agent
// //
// SELECT // SELECT conversation_id as session_id,
// conversation_id as session_id,
// agent_id, // agent_id,
// SUM(COALESCE(tokens_used, 0)) as total_tokens, // SUM(COALESCE(tokens_used, 0)) as total_tokens,
// COUNT(*) as task_count // COUNT(*) as task_count
// FROM task_completions // FROM task_completions
// WHERE created_at > strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-24 hours') // WHERE created_at > strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-24 hours')
// GROUP BY conversation_id, agent_id // GROUP BY conversation_id,
// agent_id
// HAVING SUM(COALESCE(tokens_used, 0)) > ?1 // HAVING SUM(COALESCE(tokens_used, 0)) > ?1
// ORDER BY total_tokens DESC // ORDER BY total_tokens DESC
// LIMIT ?2 // LIMIT ?2
@ -550,6 +551,14 @@ type Querier interface {
// AND key = ?2 // AND key = ?2
// LIMIT 1 // LIMIT 1
GetKV(ctx context.Context, arg GetKVParams) (AgentKv, error) GetKV(ctx context.Context, arg GetKVParams) (AgentKv, error)
//GetLatestAgentConversationByTitle
//
// SELECT id, title, created_at, updated_at
// FROM agent_conversations
// WHERE title = ?
// ORDER BY created_at DESC
// LIMIT 1
GetLatestAgentConversationByTitle(ctx context.Context, arg GetLatestAgentConversationByTitleParams) (AgentConversation, error)
//GetLatestAgentRunByConversationID //GetLatestAgentRunByConversationID
// //
// SELECT id, conversation_id, status, metadata_json, created_at, updated_at // SELECT id, conversation_id, status, metadata_json, created_at, updated_at
@ -711,7 +720,6 @@ type Querier interface {
// Task baseline queries for per-agent performance statistics // Task baseline queries for per-agent performance statistics
// Get the baseline statistics for an agent // Get the baseline statistics for an agent
// //
//
// SELECT agent_id, // SELECT agent_id,
// count, // count,
// mean_tokens, // mean_tokens,
@ -792,8 +800,11 @@ type Querier interface {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// ) // )
@ -806,6 +817,9 @@ type Querier interface {
// ?6, // ?6,
// ?7, // ?7,
// ?8, // ?8,
// ?9,
// ?10,
// ?11,
// strftime('%Y-%m-%dT%H:%M:%fZ', 'now') // strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
// ) // )
// RETURNING id, // RETURNING id,
@ -813,8 +827,11 @@ type Querier interface {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) (AgentAuditLog, error) InsertAuditEntry(ctx context.Context, arg InsertAuditEntryParams) (AgentAuditLog, error)
@ -1338,8 +1355,11 @@ type Querier interface {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -1354,8 +1374,11 @@ type Querier interface {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -1371,8 +1394,11 @@ type Querier interface {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -1388,8 +1414,11 @@ type Querier interface {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
@ -1403,13 +1432,17 @@ type Querier interface {
// session_key, // session_key,
// action, // action,
// target, // target,
// tool_call_id,
// input, // input,
// output, // output,
// success,
// error_msg,
// duration_ms, // duration_ms,
// created_at // created_at
// FROM agent_audit_log // FROM agent_audit_log
// WHERE julianday(created_at) > julianday(?1) // WHERE julianday(created_at) > julianday(?1)
// ORDER BY created_at ASC, id ASC // ORDER BY created_at ASC,
// id ASC
// LIMIT ?3 OFFSET ?2 // LIMIT ?3 OFFSET ?2
ListAuditEntriesGlobalSincePaged(ctx context.Context, arg ListAuditEntriesGlobalSincePagedParams) ([]AgentAuditLog, error) ListAuditEntriesGlobalSincePaged(ctx context.Context, arg ListAuditEntriesGlobalSincePagedParams) ([]AgentAuditLog, error)
//ListDAGEdgesBySnapshotID //ListDAGEdgesBySnapshotID
@ -1682,6 +1715,7 @@ type Querier interface {
// WHERE agent_id = ?1 // WHERE agent_id = ?1
// AND session_key = ?2 // AND session_key = ?2
// AND tags = 'session-message' // AND tags = 'session-message'
// AND suppressed_at IS NULL
// AND ( // AND (
// role = ?3 // role = ?3
// OR ?3 = '' // OR ?3 = ''
@ -1707,6 +1741,7 @@ type Querier interface {
// WHERE agent_id = ?1 // WHERE agent_id = ?1
// AND session_key = ?2 // AND session_key = ?2
// AND tags = 'session-message' // AND tags = 'session-message'
// AND suppressed_at IS NULL
// AND ( // AND (
// role = ?3 // role = ?3
// OR ?3 = '' // OR ?3 = ''
@ -1889,9 +1924,14 @@ type Querier interface {
// Store a memory retrieval record for a task // Store a memory retrieval record for a task
// //
// INSERT INTO task_retrievals (id, task_id, memory_id, similarity) // INSERT INTO task_retrievals (id, task_id, memory_id, similarity)
// VALUES (?1, ?2, ?3, ?4) // VALUES (
// ON CONFLICT (task_id, memory_id) DO UPDATE SET // ?1,
// similarity = excluded.similarity // ?2,
// ?3,
// ?4
// ) ON CONFLICT (task_id, memory_id) DO
// UPDATE
// SET similarity = excluded.similarity
StoreTaskRetrieval(ctx context.Context, arg StoreTaskRetrievalParams) error StoreTaskRetrieval(ctx context.Context, arg StoreTaskRetrievalParams) error
//UpdateAgentConversationTitle //UpdateAgentConversationTitle
// //
@ -1991,8 +2031,7 @@ type Querier interface {
// ?7, // ?7,
// ?8, // ?8,
// datetime('now') // datetime('now')
// ) // ) ON CONFLICT (agent_id) DO
// ON CONFLICT (agent_id) DO
// UPDATE // UPDATE
// SET count = excluded.count, // SET count = excluded.count,
// mean_tokens = excluded.mean_tokens, // mean_tokens = excluded.mean_tokens,

View file

@ -6,8 +6,11 @@ INSERT INTO agent_audit_log (
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
) )
@ -17,8 +20,11 @@ VALUES (
sqlc.arg(session_key), sqlc.arg(session_key),
sqlc.arg(action), sqlc.arg(action),
sqlc.arg(target), sqlc.arg(target),
sqlc.arg(tool_call_id),
sqlc.arg(input), sqlc.arg(input),
sqlc.arg(output), sqlc.arg(output),
sqlc.arg(success),
sqlc.arg(error_msg),
sqlc.arg(duration_ms), sqlc.arg(duration_ms),
strftime('%Y-%m-%dT%H:%M:%fZ', 'now') strftime('%Y-%m-%dT%H:%M:%fZ', 'now')
) )
@ -27,8 +33,11 @@ RETURNING id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at; created_at;
-- name: ListAuditEntries :many -- name: ListAuditEntries :many
@ -37,8 +46,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -51,8 +63,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -64,8 +79,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -79,8 +97,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log
@ -98,8 +119,11 @@ SELECT id,
session_key, session_key,
action, action,
target, target,
tool_call_id,
input, input,
output, output,
success,
error_msg,
duration_ms, duration_ms,
created_at created_at
FROM agent_audit_log FROM agent_audit_log

View file

@ -7,6 +7,12 @@ SELECT *
FROM agent_conversations FROM agent_conversations
WHERE id = ? WHERE id = ?
LIMIT 1; LIMIT 1;
-- name: GetLatestAgentConversationByTitle :one
SELECT *
FROM agent_conversations
WHERE title = ?
ORDER BY created_at DESC
LIMIT 1;
-- name: ListAgentConversations :many -- name: ListAgentConversations :many
SELECT * SELECT *
FROM agent_conversations FROM agent_conversations

View file

@ -219,6 +219,7 @@ FROM recall_items
WHERE agent_id = sqlc.arg(agent_id) WHERE agent_id = sqlc.arg(agent_id)
AND session_key = sqlc.arg(session_key) AND session_key = sqlc.arg(session_key)
AND tags = 'session-message' AND tags = 'session-message'
AND suppressed_at IS NULL
AND ( AND (
role = sqlc.arg(role) role = sqlc.arg(role)
OR sqlc.arg(role) = '' OR sqlc.arg(role) = ''
@ -242,6 +243,7 @@ FROM recall_items
WHERE agent_id = sqlc.arg(agent_id) WHERE agent_id = sqlc.arg(agent_id)
AND session_key = sqlc.arg(session_key) AND session_key = sqlc.arg(session_key)
AND tags = 'session-message' AND tags = 'session-message'
AND suppressed_at IS NULL
AND ( AND (
role = sqlc.arg(role) role = sqlc.arg(role)
OR sqlc.arg(role) = '' OR sqlc.arg(role) = ''
@ -253,4 +255,5 @@ SELECT COUNT(*)
FROM recall_items FROM recall_items
WHERE agent_id = sqlc.arg(agent_id) WHERE agent_id = sqlc.arg(agent_id)
AND session_key = sqlc.arg(session_key) AND session_key = sqlc.arg(session_key)
AND tags = 'session-message'; AND tags = 'session-message'
AND suppressed_at IS NULL;

View file

@ -53,6 +53,7 @@ FROM recall_items
WHERE agent_id = ?1 WHERE agent_id = ?1
AND session_key = ?2 AND session_key = ?2
AND tags = 'session-message' AND tags = 'session-message'
AND suppressed_at IS NULL
` `
type CountSessionMessagesParams struct { type CountSessionMessagesParams struct {
@ -67,6 +68,7 @@ type CountSessionMessagesParams struct {
// WHERE agent_id = ?1 // WHERE agent_id = ?1
// AND session_key = ?2 // AND session_key = ?2
// AND tags = 'session-message' // AND tags = 'session-message'
// AND suppressed_at IS NULL
func (q *Queries) CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error) { func (q *Queries) CountSessionMessages(ctx context.Context, arg CountSessionMessagesParams) (int64, error) {
row := q.db.QueryRowContext(ctx, CountSessionMessages, arg.AgentID, arg.SessionKey) row := q.db.QueryRowContext(ctx, CountSessionMessages, arg.AgentID, arg.SessionKey)
var count int64 var count int64
@ -769,6 +771,7 @@ FROM recall_items
WHERE agent_id = ?1 WHERE agent_id = ?1
AND session_key = ?2 AND session_key = ?2
AND tags = 'session-message' AND tags = 'session-message'
AND suppressed_at IS NULL
AND ( AND (
role = ?3 role = ?3
OR ?3 = '' OR ?3 = ''
@ -817,6 +820,7 @@ type ListSessionMessagesRow struct {
// WHERE agent_id = ?1 // WHERE agent_id = ?1
// AND session_key = ?2 // AND session_key = ?2
// AND tags = 'session-message' // AND tags = 'session-message'
// AND suppressed_at IS NULL
// AND ( // AND (
// role = ?3 // role = ?3
// OR ?3 = '' // OR ?3 = ''
@ -881,6 +885,7 @@ FROM recall_items
WHERE agent_id = ?1 WHERE agent_id = ?1
AND session_key = ?2 AND session_key = ?2
AND tags = 'session-message' AND tags = 'session-message'
AND suppressed_at IS NULL
AND ( AND (
role = ?3 role = ?3
OR ?3 = '' OR ?3 = ''
@ -930,6 +935,7 @@ type ListSessionMessagesPagedRow struct {
// WHERE agent_id = ?1 // WHERE agent_id = ?1
// AND session_key = ?2 // AND session_key = ?2
// AND tags = 'session-message' // AND tags = 'session-message'
// AND suppressed_at IS NULL
// AND ( // AND (
// role = ?3 // role = ?3
// OR ?3 = '' // OR ?3 = ''

View file

@ -91,14 +91,14 @@ func (q *Queries) GetCompletedTasks(ctx context.Context, arg GetCompletedTasksPa
} }
const GetHighTokenSessions = `-- name: GetHighTokenSessions :many const GetHighTokenSessions = `-- name: GetHighTokenSessions :many
SELECT SELECT conversation_id as session_id,
conversation_id as session_id,
agent_id, agent_id,
SUM(COALESCE(tokens_used, 0)) as total_tokens, SUM(COALESCE(tokens_used, 0)) as total_tokens,
COUNT(*) as task_count COUNT(*) as task_count
FROM task_completions FROM task_completions
WHERE created_at > strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-24 hours') WHERE created_at > strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-24 hours')
GROUP BY conversation_id, agent_id GROUP BY conversation_id,
agent_id
HAVING SUM(COALESCE(tokens_used, 0)) > ?1 HAVING SUM(COALESCE(tokens_used, 0)) > ?1
ORDER BY total_tokens DESC ORDER BY total_tokens DESC
LIMIT ?2 LIMIT ?2
@ -118,14 +118,14 @@ type GetHighTokenSessionsRow struct {
// Get sessions with high token usage grouped by conversation/agent // Get sessions with high token usage grouped by conversation/agent
// //
// SELECT // SELECT conversation_id as session_id,
// conversation_id as session_id,
// agent_id, // agent_id,
// SUM(COALESCE(tokens_used, 0)) as total_tokens, // SUM(COALESCE(tokens_used, 0)) as total_tokens,
// COUNT(*) as task_count // COUNT(*) as task_count
// FROM task_completions // FROM task_completions
// WHERE created_at > strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-24 hours') // WHERE created_at > strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-24 hours')
// GROUP BY conversation_id, agent_id // GROUP BY conversation_id,
// agent_id
// HAVING SUM(COALESCE(tokens_used, 0)) > ?1 // HAVING SUM(COALESCE(tokens_used, 0)) > ?1
// ORDER BY total_tokens DESC // ORDER BY total_tokens DESC
// LIMIT ?2 // LIMIT ?2
@ -319,7 +319,6 @@ func (q *Queries) GetRetrievedMemories(ctx context.Context, arg GetRetrievedMemo
} }
const GetTaskBaseline = `-- name: GetTaskBaseline :one const GetTaskBaseline = `-- name: GetTaskBaseline :one
SELECT agent_id, SELECT agent_id,
count, count,
mean_tokens, mean_tokens,
@ -682,9 +681,14 @@ func (q *Queries) StoreTaskCompletion(ctx context.Context, arg StoreTaskCompleti
const StoreTaskRetrieval = `-- name: StoreTaskRetrieval :exec const StoreTaskRetrieval = `-- name: StoreTaskRetrieval :exec
INSERT INTO task_retrievals (id, task_id, memory_id, similarity) INSERT INTO task_retrievals (id, task_id, memory_id, similarity)
VALUES (?1, ?2, ?3, ?4) VALUES (
ON CONFLICT (task_id, memory_id) DO UPDATE SET ?1,
similarity = excluded.similarity ?2,
?3,
?4
) ON CONFLICT (task_id, memory_id) DO
UPDATE
SET similarity = excluded.similarity
` `
type StoreTaskRetrievalParams struct { type StoreTaskRetrievalParams struct {
@ -697,9 +701,14 @@ type StoreTaskRetrievalParams struct {
// Store a memory retrieval record for a task // Store a memory retrieval record for a task
// //
// INSERT INTO task_retrievals (id, task_id, memory_id, similarity) // INSERT INTO task_retrievals (id, task_id, memory_id, similarity)
// VALUES (?1, ?2, ?3, ?4) // VALUES (
// ON CONFLICT (task_id, memory_id) DO UPDATE SET // ?1,
// similarity = excluded.similarity // ?2,
// ?3,
// ?4
// ) ON CONFLICT (task_id, memory_id) DO
// UPDATE
// SET similarity = excluded.similarity
func (q *Queries) StoreTaskRetrieval(ctx context.Context, arg StoreTaskRetrievalParams) error { func (q *Queries) StoreTaskRetrieval(ctx context.Context, arg StoreTaskRetrievalParams) error {
_, err := q.db.ExecContext(ctx, StoreTaskRetrieval, _, err := q.db.ExecContext(ctx, StoreTaskRetrieval,
arg.ID, arg.ID,
@ -792,8 +801,7 @@ VALUES (
?7, ?7,
?8, ?8,
datetime('now') datetime('now')
) ) ON CONFLICT (agent_id) DO
ON CONFLICT (agent_id) DO
UPDATE UPDATE
SET count = excluded.count, SET count = excluded.count,
mean_tokens = excluded.mean_tokens, mean_tokens = excluded.mean_tokens,
@ -839,8 +847,7 @@ type UpdateTaskBaselineParams struct {
// ?7, // ?7,
// ?8, // ?8,
// datetime('now') // datetime('now')
// ) // ) ON CONFLICT (agent_id) DO
// ON CONFLICT (agent_id) DO
// UPDATE // UPDATE
// SET count = excluded.count, // SET count = excluded.count,
// mean_tokens = excluded.mean_tokens, // mean_tokens = excluded.mean_tokens,

View file

@ -27,11 +27,15 @@ CREATE TABLE IF NOT EXISTS recall_items (
tags TEXT NOT NULL DEFAULT '', tags TEXT NOT NULL DEFAULT '',
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
suppressed_at DATETIME, -- soft delete timestamp (T2.4) suppressed_at DATETIME,
-- soft delete timestamp (T2.4)
-- RL (Memelord) support columns -- RL (Memelord) support columns
rl_weight REAL DEFAULT 1.0, -- current weight for credit assignment rl_weight REAL DEFAULT 1.0,
rl_credit REAL, -- accumulated credit for this memory -- current weight for credit assignment
self_report_score INTEGER, -- self-reported usefulness score rl_credit REAL,
-- accumulated credit for this memory
self_report_score INTEGER,
-- self-reported usefulness score
task_retrieval_count INTEGER DEFAULT 0 -- how many times retrieved for tasks task_retrieval_count INTEGER DEFAULT 0 -- how many times retrieved for tasks
); );
CREATE INDEX IF NOT EXISTS idx_recall_agent_session ON recall_items(agent_id, session_key); CREATE INDEX IF NOT EXISTS idx_recall_agent_session ON recall_items(agent_id, session_key);
@ -96,13 +100,17 @@ CREATE TABLE IF NOT EXISTS agent_audit_log (
session_key TEXT NOT NULL DEFAULT '', session_key TEXT NOT NULL DEFAULT '',
action TEXT NOT NULL, action TEXT NOT NULL,
target TEXT NOT NULL DEFAULT '', target TEXT NOT NULL DEFAULT '',
tool_call_id TEXT NOT NULL DEFAULT '',
input TEXT, input TEXT,
output TEXT, output TEXT,
success BOOLEAN NOT NULL DEFAULT TRUE,
error_msg TEXT NOT NULL DEFAULT '',
duration_ms INTEGER, duration_ms INTEGER,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
); );
CREATE INDEX IF NOT EXISTS idx_audit_agent_time ON agent_audit_log(agent_id, created_at DESC); CREATE INDEX IF NOT EXISTS idx_audit_agent_time ON agent_audit_log(agent_id, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_audit_action ON agent_audit_log(action); CREATE INDEX IF NOT EXISTS idx_audit_action ON agent_audit_log(action);
CREATE INDEX IF NOT EXISTS idx_audit_tool_call_id ON agent_audit_log(tool_call_id);
-- ============================================================================ -- ============================================================================
-- Agent Runtime State Tables -- Agent Runtime State Tables
-- ============================================================================ -- ============================================================================
@ -422,7 +430,6 @@ CREATE TABLE IF NOT EXISTS memory_edges (
CREATE INDEX IF NOT EXISTS idx_memory_edges_from ON memory_edges(from_id); CREATE INDEX IF NOT EXISTS idx_memory_edges_from ON memory_edges(from_id);
CREATE INDEX IF NOT EXISTS idx_memory_edges_to ON memory_edges(to_id); CREATE INDEX IF NOT EXISTS idx_memory_edges_to ON memory_edges(to_id);
CREATE INDEX IF NOT EXISTS idx_memory_edges_type ON memory_edges(edge_type); CREATE INDEX IF NOT EXISTS idx_memory_edges_type ON memory_edges(edge_type);
-- ============================================================================ -- ============================================================================
-- RL (Reinforcement Learning) Support Tables -- RL (Reinforcement Learning) Support Tables
-- ============================================================================ -- ============================================================================
@ -438,7 +445,6 @@ CREATE TABLE IF NOT EXISTS task_baselines (
m2_user_corrections REAL DEFAULT 0, m2_user_corrections REAL DEFAULT 0,
updated_at DATETIME updated_at DATETIME
); );
-- Task completions: record of completed agent runs for RL analysis -- Task completions: record of completed agent runs for RL analysis
CREATE TABLE IF NOT EXISTS task_completions ( CREATE TABLE IF NOT EXISTS task_completions (
id BLOB PRIMARY KEY, id BLOB PRIMARY KEY,
@ -455,7 +461,6 @@ CREATE TABLE IF NOT EXISTS task_completions (
); );
CREATE INDEX IF NOT EXISTS idx_task_completions_agent_created ON task_completions(agent_id, created_at); CREATE INDEX IF NOT EXISTS idx_task_completions_agent_created ON task_completions(agent_id, created_at);
CREATE INDEX IF NOT EXISTS idx_task_completions_run ON task_completions(run_id); CREATE INDEX IF NOT EXISTS idx_task_completions_run ON task_completions(run_id);
-- Task retrievals: links memories retrieved during task execution for RL credit assignment -- Task retrievals: links memories retrieved during task execution for RL credit assignment
CREATE TABLE IF NOT EXISTS task_retrievals ( CREATE TABLE IF NOT EXISTS task_retrievals (
id BLOB PRIMARY KEY, id BLOB PRIMARY KEY,

View file

@ -336,23 +336,31 @@ func (m *MemoryStore) Search(ctx context.Context, query string, opts memory.Sear
var baselineSets [][]memory.SearchResult var baselineSets [][]memory.SearchResult
var baselineWeights []float64 var baselineWeights []float64
// 1. Keyword search (via delegate)
kwWeight := opts.KeywordWeight kwWeight := opts.KeywordWeight
if kwWeight <= 0 { vecWeight := opts.VectorWeight
if kwWeight == 0 && vecWeight == 0 {
kwWeight = 1.0 kwWeight = 1.0
vecWeight = 0.8
} }
kwResults, err := m.keywordSearch(ctx, query, opts, limit*2) // fetch extra for fusion if kwWeight < 0 {
if err == nil && len(kwResults) > 0 { kwWeight = 0
baselineSets = append(baselineSets, kwResults) }
baselineWeights = append(baselineWeights, kwWeight)
if vecWeight < 0 {
vecWeight = 0
}
// 1. Keyword search (via delegate)
if kwWeight > 0 {
kwResults, err := m.keywordSearch(ctx, query, opts, limit*2) // fetch extra for fusion
if err == nil && len(kwResults) > 0 {
baselineSets = append(baselineSets, kwResults)
baselineWeights = append(baselineWeights, kwWeight)
}
} }
// 2. Vector search (if embedder available) // 2. Vector search (if embedder available)
vecWeight := opts.VectorWeight if vecWeight > 0 && m.embedder != nil {
if vecWeight <= 0 {
vecWeight = 0.8
}
if m.embedder != nil {
vecResults, err := m.vectorSearch(ctx, query, opts, limit*2) vecResults, err := m.vectorSearch(ctx, query, opts, limit*2)
if err == nil && len(vecResults) > 0 { if err == nil && len(vecResults) > 0 {
baselineSets = append(baselineSets, vecResults) baselineSets = append(baselineSets, vecResults)
@ -481,11 +489,10 @@ func (m *MemoryStore) hybridProjectionSearch(ctx context.Context, query string,
// Working-context view // Working-context view
wc, err := m.delegate.GetWorkingContext(ctx, m.agentID, opts.SessionKey) wc, err := m.delegate.GetWorkingContext(ctx, m.agentID, opts.SessionKey)
if err == nil && wc != nil && strings.TrimSpace(wc.Content) != "" { if err == nil && wc != nil && strings.TrimSpace(wc.Content) != "" &&
score := 0.4 queryLower != "" &&
if queryLower != "" && strings.Contains(strings.ToLower(wc.Content), queryLower) { strings.Contains(strings.ToLower(wc.Content), queryLower) &&
score = 0.95 !looksLikeMemorySearchPromptEcho(wc.Content, query) {
}
content := wc.Content content := wc.Content
if len(content) > 1200 { if len(content) > 1200 {
content = content[:1200] + "..." content = content[:1200] + "..."
@ -494,7 +501,7 @@ func (m *MemoryStore) hybridProjectionSearch(ctx context.Context, query string,
ID: ids.New(), ID: ids.New(),
Content: content, Content: content,
Source: "working-context:" + opts.SessionKey, Source: "working-context:" + opts.SessionKey,
Score: score, Score: 0.95,
Sector: memory.SectorReflective, Sector: memory.SectorReflective,
}) })
} }
@ -517,6 +524,9 @@ func (m *MemoryStore) hybridProjectionSearch(ctx context.Context, query string,
if queryLower != "" && !strings.Contains(strings.ToLower(node.Summary), queryLower) { if queryLower != "" && !strings.Contains(strings.ToLower(node.Summary), queryLower) {
continue continue
} }
if looksLikeMemorySearchPromptEcho(node.Summary, query) {
continue
}
score := 0.55 + (0.1 * float64(node.Level)) score := 0.55 + (0.1 * float64(node.Level))
if score > 0.95 { if score > 0.95 {
score = 0.95 score = 0.95
@ -608,7 +618,10 @@ func (m *MemoryStore) keywordSearch(ctx context.Context, query string, opts memo
if m.delegate.HasFTS() { if m.delegate.HasFTS() {
items, err := m.delegate.SearchRecallByFTS(ctx, query, opts.AgentID, limit) items, err := m.delegate.SearchRecallByFTS(ctx, query, opts.AgentID, limit)
if err == nil && len(items) > 0 { if err == nil && len(items) > 0 {
return recallItemsToResults(items), nil items = filterSearchPromptEchoItems(items, query)
if len(items) > 0 {
return recallItemsToResults(items), nil
}
} }
} }
@ -617,6 +630,7 @@ func (m *MemoryStore) keywordSearch(ctx context.Context, query string, opts memo
if err != nil { if err != nil {
return nil, err return nil, err
} }
items = filterSearchPromptEchoItems(items, query)
return recallItemsToResults(items), nil return recallItemsToResults(items), nil
} }
@ -740,6 +754,60 @@ func recallItemsToResults(items []*memory.RecallItem) []memory.SearchResult {
return results return results
} }
func filterSearchPromptEchoItems(items []*memory.RecallItem, query string) []*memory.RecallItem {
if len(items) == 0 {
return nil
}
filtered := items[:0]
for _, item := range items {
if shouldSuppressSearchPromptEcho(item, query) {
continue
}
filtered = append(filtered, item)
}
return filtered
}
func shouldSuppressSearchPromptEcho(item *memory.RecallItem, query string) bool {
if item == nil {
return false
}
if item.Role != "user" {
return false
}
if !strings.Contains(strings.ToLower(item.Tags), "session-message") {
return false
}
return looksLikeMemorySearchPromptEcho(item.Content, query)
}
func looksLikeMemorySearchPromptEcho(content, query string) bool {
lowerContent := strings.ToLower(strings.TrimSpace(content))
lowerQuery := strings.ToLower(strings.TrimSpace(query))
if lowerContent == "" || lowerQuery == "" {
return false
}
if !strings.Contains(lowerContent, lowerQuery) {
return false
}
for _, needle := range []string{
"search your memory",
"search my memory",
"look in your memory",
"look through your memory",
"what do you remember",
"tell me what you find",
"tell me what you remember",
} {
if strings.Contains(lowerContent, needle) {
return true
}
}
return false
}
// --- Summaries --- // --- Summaries ---
func (m *MemoryStore) StoreSummary(ctx context.Context, summary *memory.MemorySummary) error { func (m *MemoryStore) StoreSummary(ctx context.Context, summary *memory.MemorySummary) error {

View file

@ -257,6 +257,47 @@ func TestSearch_KeywordOnly(t *testing.T) {
assert.Contains(t, results[0].Content, "generics") assert.Contains(t, results[0].Content, "generics")
} }
func TestSearch_KeywordOnlySkipsVectorWhenWeightZero(t *testing.T) {
t.Parallel()
ctx := t.Context()
del, err := delegate.NewLibSQLInMemory()
require.NoError(t, err)
require.NoError(t, del.Init(ctx))
emb := &countingEmbedder{dims: 768}
store := New(del, NewMarkdownChunker(MarkdownChunkerConfig{
ChunkSize: 200,
ChunkOverlap: 40,
}), emb, Config{
ContextWindowTokens: 10000,
OffloadThresholdTokens: 100,
DefaultHalfLifeHours: 168,
})
store.SetAgentID("agent-1")
t.Cleanup(func() { store.Close() })
require.NoError(t, store.StoreRecall(ctx, &memory.RecallItem{
AgentID: "agent-1",
SessionKey: "s1",
Role: "user",
Sector: memory.SectorSemantic,
Importance: 0.9,
Content: "Go generics were introduced in Go 1.18",
}))
results, err := store.Search(ctx, "generics", memory.SearchOptions{
AgentID: "agent-1",
Limit: 10,
KeywordWeight: 1.0,
VectorWeight: 0.0,
})
require.NoError(t, err)
assert.NotEmpty(t, results)
assert.Zero(t, emb.embedCalls.Load(), "keyword-only search should not invoke vector embedding")
assert.Zero(t, emb.batchCalls.Load(), "keyword-only search should not invoke batch embeddings")
}
func TestSearch_HybridWithEmbeddings(t *testing.T) { func TestSearch_HybridWithEmbeddings(t *testing.T) {
t.Parallel() t.Parallel()
ctx := t.Context() ctx := t.Context()

View file

@ -3,9 +3,10 @@ package store
import ( import (
"context" "context"
"fmt" "fmt"
jsonv2 "github.com/go-json-experiment/json"
"strings" "strings"
jsonv2 "github.com/go-json-experiment/json"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids" "github.com/ZanzyTHEbar/dragonscale/pkg/ids"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory" "github.com/ZanzyTHEbar/dragonscale/pkg/memory"
) )
@ -131,9 +132,10 @@ func (t *MemoryTool) search(ctx context.Context, req *MemoryToolRequest) (*Memor
} }
results, err := t.store.Search(ctx, req.Query, memory.SearchOptions{ results, err := t.store.Search(ctx, req.Query, memory.SearchOptions{
AgentID: t.agentID, AgentID: t.agentID,
Sectors: sectors, SessionKey: t.session,
Limit: limit, Sectors: sectors,
Limit: limit,
}) })
if err != nil { if err != nil {
return nil, err return nil, err
@ -152,11 +154,18 @@ func (t *MemoryTool) search(ctx context.Context, req *MemoryToolRequest) (*Memor
return &MemoryToolResponse{ return &MemoryToolResponse{
Success: true, Success: true,
Message: fmt.Sprintf("Found %d results for: %s", len(entries), req.Query), Message: searchSummaryMessage(len(entries), req.Query),
Results: entries, Results: entries,
}, nil }, nil
} }
func searchSummaryMessage(count int, query string) string {
if count == 0 {
return fmt.Sprintf("No results found for: %s", query)
}
return fmt.Sprintf("Found %d results for: %s", count, query)
}
func (t *MemoryTool) read(ctx context.Context, req *MemoryToolRequest) (*MemoryToolResponse, error) { func (t *MemoryTool) read(ctx context.Context, req *MemoryToolRequest) (*MemoryToolResponse, error) {
if req.ID == "" { if req.ID == "" {
return &MemoryToolResponse{Success: false, Message: "id is required for read"}, nil return &MemoryToolResponse{Success: false, Message: "id is required for read"}, nil

View file

@ -1,10 +1,12 @@
package store package store
import ( import (
"strings"
"testing" "testing"
jsonv2 "github.com/go-json-experiment/json" jsonv2 "github.com/go-json-experiment/json"
"github.com/ZanzyTHEbar/dragonscale/pkg/memory"
"github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
@ -77,6 +79,74 @@ func TestMemoryTool_Search(t *testing.T) {
assert.Contains(t, resp.Results[0].Content, "goroutines") assert.Contains(t, resp.Results[0].Content, "goroutines")
} }
func TestMemoryTool_SearchNoResultsUsesExplicitEmptyMessage(t *testing.T) {
t.Parallel()
tool := newTestMemoryTool(t)
resp := executeAndParse(t, tool, `{"action":"search","query":"xyzzy_nonexistent_topic_42","limit":5}`)
assert.True(t, resp.Success)
assert.Empty(t, resp.Results)
assert.Contains(t, resp.Message, "No results found for:")
}
func TestMemoryTool_SearchSuppressesPromptEchoArtifacts(t *testing.T) {
t.Parallel()
tool := newTestMemoryTool(t)
ctx := t.Context()
promptEcho := "Search your memory for 'xyzzy_nonexistent_topic_42' and tell me what you find."
require.NoError(t, tool.store.StoreRecall(ctx, &memory.RecallItem{
AgentID: "agent-1",
SessionKey: "session-1",
Role: "user",
Sector: memory.SectorEpisodic,
Importance: 0.5,
Content: promptEcho,
Tags: "session-message",
}))
require.NoError(t, tool.store.SetWorkingContext(ctx, "agent-1", "session-1", promptEcho))
resp := executeAndParse(t, tool, `{"action":"search","query":"xyzzy_nonexistent_topic_42","limit":5}`)
assert.True(t, resp.Success)
assert.Empty(t, resp.Results)
assert.Contains(t, resp.Message, "No results found for:")
}
func TestMemoryTool_SearchKeepsRealMemoryWhenPromptEchoExists(t *testing.T) {
t.Parallel()
tool := newTestMemoryTool(t)
ctx := t.Context()
promptEcho := "Search your memory for 'xyzzy_nonexistent_topic_42' and tell me what you find."
require.NoError(t, tool.store.StoreRecall(ctx, &memory.RecallItem{
AgentID: "agent-1",
SessionKey: "session-1",
Role: "user",
Sector: memory.SectorEpisodic,
Importance: 0.5,
Content: promptEcho,
Tags: "session-message",
}))
require.NoError(t, tool.store.StoreRecall(ctx, &memory.RecallItem{
AgentID: "agent-1",
SessionKey: "session-1",
Role: "assistant",
Sector: memory.SectorSemantic,
Importance: 0.9,
Content: "Regression ledger entry: xyzzy_nonexistent_topic_42 was intentionally left undefined.",
Tags: "semantic-fact",
}))
require.NoError(t, tool.store.SetWorkingContext(ctx, "agent-1", "session-1", promptEcho))
resp := executeAndParse(t, tool, `{"action":"search","query":"xyzzy_nonexistent_topic_42","limit":5}`)
assert.True(t, resp.Success)
require.NotEmpty(t, resp.Results)
assert.Contains(t, resp.Results[0].Content, "intentionally left undefined")
for _, entry := range resp.Results {
assert.False(t, strings.Contains(strings.ToLower(entry.Content), "search your memory"))
}
}
func TestMemoryTool_Update(t *testing.T) { func TestMemoryTool_Update(t *testing.T) {
t.Parallel() t.Parallel()
tool := newTestMemoryTool(t) tool := newTestMemoryTool(t)

View file

@ -0,0 +1,26 @@
package memory
import (
"time"
"github.com/ZanzyTHEbar/dragonscale/pkg/ids"
)
// MemoryRating represents a self-reported usefulness score for a memory.
type MemoryRating struct {
MemoryID ids.UUID
Score int // 0-3 scale
}
// TaskCompletionRecord tracks the outcome of an agent run for RL analysis.
type TaskCompletionRecord struct {
TaskID string
Description string
TokensUsed int
ToolCalls int
Errors int
UserCorrections int
Completed bool
SelfReports []MemoryRating
CreatedAt time.Time
}

View file

@ -57,15 +57,17 @@ type msgPersistItem struct {
} }
type SessionManager struct { type SessionManager struct {
sessions map[string]*Session // primary store (always authoritative) sessions map[string]*Session // primary store (always authoritative)
lru *cache.LRU[string, bool] // tracks access order; value is just a presence flag lru *cache.LRU[string, bool] // tracks access order; value is just a presence flag
mu sync.RWMutex mu sync.RWMutex
storage string msgMu sync.RWMutex
cfg SessionManagerConfig storage string
delegate memory.MemoryDelegate cfg SessionManagerConfig
agentID string delegate memory.MemoryDelegate
msgChan chan msgPersistItem // async message persistence; nil when delegate is nil agentID string
msgDone chan struct{} // closed when the persist worker exits msgChan chan msgPersistItem // async message persistence; nil when delegate is nil
msgDone chan struct{} // closed when the persist worker exits
msgClosed bool
} }
func NewSessionManager(storage string, opts ...SessionOption) *SessionManager { func NewSessionManager(storage string, opts ...SessionOption) *SessionManager {
@ -223,6 +225,16 @@ func (sm *SessionManager) loadSessionsFromDelegate() {
} }
} }
if summaries, err := sm.delegate.ListSummaries(ctx, sm.agentID, sessionKey, 1); err == nil && len(summaries) > 0 {
session.Summary = summaries[0].Content
if summaries[0].CreatedAt.After(session.Updated) {
session.Updated = summaries[0].CreatedAt
}
if session.Created.IsZero() || summaries[0].CreatedAt.Before(session.Created) {
session.Created = summaries[0].CreatedAt
}
}
// Integrity validation and projection pointer persistence for deterministic resume. // Integrity validation and projection pointer persistence for deterministic resume.
restoredPtr := projectFromItems(chronItems) restoredPtr := projectFromItems(chronItems)
if restoredPtr != nil { if restoredPtr != nil {
@ -379,15 +391,19 @@ func (sm *SessionManager) msgPersistWorker() {
} }
func (sm *SessionManager) enqueuePersistItem(item msgPersistItem) (ok bool) { func (sm *SessionManager) enqueuePersistItem(item msgPersistItem) (ok bool) {
if sm.msgChan == nil { sm.msgMu.RLock()
if sm.msgChan == nil || sm.msgClosed {
sm.msgMu.RUnlock()
return false return false
} }
ch := sm.msgChan
defer func() { defer func() {
sm.msgMu.RUnlock()
if recover() != nil { if recover() != nil {
ok = false ok = false
} }
}() }()
sm.msgChan <- item ch <- item
return true return true
} }
@ -602,15 +618,34 @@ func (sm *SessionManager) Flush() {
// Close drains the async message persistence channel and waits for completion. // Close drains the async message persistence channel and waits for completion.
func (sm *SessionManager) Close() { func (sm *SessionManager) Close() {
if sm.msgChan != nil { sm.msgMu.Lock()
close(sm.msgChan) if sm.msgChan == nil || sm.msgClosed {
<-sm.msgDone sm.msgMu.Unlock()
return
} }
sm.msgClosed = true
ch := sm.msgChan
done := sm.msgDone
close(ch)
sm.msgMu.Unlock()
<-done
} }
func (sm *SessionManager) Save(key string) error { func (sm *SessionManager) Save(key string) error {
if sm.delegate != nil { if sm.delegate != nil {
return nil sm.Flush()
sm.mu.RLock()
stored, ok := sm.sessions[key]
if !ok {
sm.mu.RUnlock()
return nil
}
snapshot := snapshotSession(stored)
sm.mu.RUnlock()
return sm.persistSummaryToDelegate(key, &snapshot)
} }
if sm.storage == "" { if sm.storage == "" {
return nil return nil
@ -640,6 +675,34 @@ func (sm *SessionManager) Save(key string) error {
return sm.writeSessionToDisk(key, &snapshot) return sm.writeSessionToDisk(key, &snapshot)
} }
func (sm *SessionManager) persistSummaryToDelegate(key string, session *Session) error {
if sm.delegate == nil || session == nil {
return nil
}
summary := strings.TrimSpace(session.Summary)
if summary == "" {
return nil
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
existing, err := sm.delegate.ListSummaries(ctx, sm.agentID, key, 1)
if err == nil && len(existing) > 0 && strings.TrimSpace(existing[0].Content) == summary {
return nil
}
return sm.delegate.InsertSummary(ctx, &memory.MemorySummary{
ID: ids.New(),
AgentID: sm.agentID,
SessionKey: key,
Content: summary,
FromMsgIdx: 0,
ToMsgIdx: len(session.Messages),
})
}
// saveSessionLocked saves a session to disk. Caller must hold sm.mu. // saveSessionLocked saves a session to disk. Caller must hold sm.mu.
func (sm *SessionManager) saveSessionLocked(key string, session *Session) { func (sm *SessionManager) saveSessionLocked(key string, session *Session) {
if sm.storage == "" { if sm.storage == "" {
@ -766,6 +829,38 @@ func (sm *SessionManager) SetHistory(key string, history []messages.Message) {
} }
} }
// ReplaceHistory atomically ensures a session exists and replaces its history
// and summary in a single critical section.
func (sm *SessionManager) ReplaceHistory(key string, history []messages.Message, summary string) {
sm.mu.Lock()
defer sm.mu.Unlock()
session, ok := sm.sessions[key]
if !ok {
if sm.storage != "" {
if loaded := sm.loadSessionFromDisk(key); loaded != nil {
session = loaded
sm.sessions[key] = session
}
}
if session == nil {
session = &Session{
Key: key,
Messages: []messages.Message{},
Created: time.Now(),
}
sm.sessions[key] = session
}
}
msgs := make([]messages.Message, len(history))
copy(msgs, history)
session.Messages = msgs
session.Summary = summary
session.Updated = time.Now()
sm.touchLRU(key)
}
func toolCallsJSON(msg messages.Message) string { func toolCallsJSON(msg messages.Message) string {
if len(msg.ToolCalls) == 0 { if len(msg.ToolCalls) == 0 {
return "" return ""

View file

@ -4,6 +4,7 @@ import (
"fmt" "fmt"
"os" "os"
"path/filepath" "path/filepath"
"sync"
"testing" "testing"
"time" "time"
@ -267,6 +268,66 @@ func TestSessionManager_DelegateSaveIsNoop(t *testing.T) {
} }
} }
func TestSessionManager_DelegateSavePersistsAndRestoresSummary(t *testing.T) {
t.Parallel()
del, err := delegate.NewLibSQLInMemory()
require.NoError(t, err)
require.NoError(t, del.Init(t.Context()))
defer del.Close()
sm := NewSessionManager("", WithSessionDelegate(del, "test-agent"))
key := "delegate-summary"
sm.AddMessage(key, "user", "hello")
sm.SetSummary(key, "durable summary")
require.NoError(t, sm.Save(key))
summaries, err := del.ListSummaries(t.Context(), "test-agent", key, 1)
require.NoError(t, err)
require.Len(t, summaries, 1)
assert.Equal(t, "durable summary", summaries[0].Content)
restored := NewSessionManager("", WithSessionDelegate(del, "test-agent"))
assert.Equal(t, "durable summary", restored.GetSummary(key))
}
func TestSessionManager_CloseWhileSaveRuns(t *testing.T) {
t.Parallel()
del, err := delegate.NewLibSQLInMemory()
require.NoError(t, err)
require.NoError(t, del.Init(t.Context()))
defer del.Close()
sm := NewSessionManager("", WithSessionDelegate(del, "test-agent"))
key := "close-while-save"
sm.AddMessage(key, "user", "hello")
sm.AddMessage(key, "assistant", "world")
var wg sync.WaitGroup
errCh := make(chan error, 8)
for i := 0; i < 8; i++ {
wg.Add(1)
go func() {
defer wg.Done()
errCh <- sm.Save(key)
}()
}
wg.Add(1)
go func() {
defer wg.Done()
sm.Close()
}()
wg.Wait()
close(errCh)
for err := range errCh {
require.NoError(t, err)
}
sm.Close()
}
func TestSessionManager_DelegateBootstrapPaginationAndOrder(t *testing.T) { func TestSessionManager_DelegateBootstrapPaginationAndOrder(t *testing.T) {
t.Parallel() t.Parallel()
del, err := delegate.NewLibSQLInMemory() del, err := delegate.NewLibSQLInMemory()

View file

@ -27,7 +27,7 @@ func (t *EditFileTool) Name() string {
} }
func (t *EditFileTool) Description() string { func (t *EditFileTool) Description() string {
return "Edit a file by replacing old_text with new_text. The old_text must exist exactly in the file." return "Edit a file by replacing old_text with new_text. The old_text must exist exactly once in the file. Example: {\"path\":\"edit_target.txt\",\"old_text\":\"world\",\"new_text\":\"dragonscale\"}."
} }
func (t *EditFileTool) Parameters() map[string]interface{} { func (t *EditFileTool) Parameters() map[string]interface{} {
@ -115,7 +115,7 @@ func (t *AppendFileTool) Name() string {
} }
func (t *AppendFileTool) Description() string { func (t *AppendFileTool) Description() string {
return "Append content to the end of a file" return "Append content to the end of a file. Example: {\"path\":\"append_test.txt\",\"content\":\"line two\\n\"}."
} }
func (t *AppendFileTool) Parameters() map[string]interface{} { func (t *AppendFileTool) Parameters() map[string]interface{} {

View file

@ -132,7 +132,7 @@ func (t *ReadFileTool) Name() string {
} }
func (t *ReadFileTool) Description() string { func (t *ReadFileTool) Description() string {
return "Read the contents of a file" return "Read the contents of a regular workspace file. Use this to verify file state after writes or edits. For skills, use skill_read instead of read_file."
} }
func (t *ReadFileTool) Parameters() map[string]interface{} { func (t *ReadFileTool) Parameters() map[string]interface{} {
@ -188,7 +188,7 @@ func (t *WriteFileTool) Name() string {
} }
func (t *WriteFileTool) Description() string { func (t *WriteFileTool) Description() string {
return "Write content to a file" return "Create or fully overwrite a file with new content. Use edit_file for targeted replacements and append_file for adding content to the end."
} }
func (t *WriteFileTool) Parameters() map[string]interface{} { func (t *WriteFileTool) Parameters() map[string]interface{} {

View file

@ -38,7 +38,7 @@ func (t *ToolSearchTool) SetFocusContext(delegate KVStore, sessionKeyFn func() s
func (t *ToolSearchTool) Name() string { return "tool_search" } func (t *ToolSearchTool) Name() string { return "tool_search" }
func (t *ToolSearchTool) Description() string { func (t *ToolSearchTool) Description() string {
return "Search for available tools and skills by keyword. Returns names, descriptions, parameter schemas, and kind (tool or skill). Discovered tools become directly callable in your next step — no need to use tool_call. For skills, use skill_read to load full content." return "Search for available tools by keyword when you do not already know the exact tool name. Returns tool names, descriptions, parameter schemas, and kind metadata. Discovered tools become directly callable in your next step — no need to use tool_call. Do NOT use tool_search for skill discovery; use skill_search instead."
} }
func (t *ToolSearchTool) Parameters() map[string]interface{} { func (t *ToolSearchTool) Parameters() map[string]interface{} {

View file

@ -10,6 +10,7 @@ import (
"path/filepath" "path/filepath"
"regexp" "regexp"
"runtime" "runtime"
"strconv"
"strings" "strings"
"syscall" "syscall"
"time" "time"
@ -148,7 +149,7 @@ func (t *ExecTool) Name() string {
} }
func (t *ExecTool) Description() string { func (t *ExecTool) Description() string {
return "Execute a shell command and return its output. Use with caution." return "Execute a shell command and return its output. Pass only the raw command string in command, for example {\"command\":\"uname -s\"}. Keep working_dir separate. Do not use exec for normal file create/edit/append tasks when write_file, edit_file, or append_file fits."
} }
func (t *ExecTool) Parameters() map[string]interface{} { func (t *ExecTool) Parameters() map[string]interface{} {
@ -185,6 +186,7 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *To
if !ok { if !ok {
return ErrorResult("command is required") return ErrorResult("command is required")
} }
command = normalizeShellCommand(command)
if len(command) > maxCommandLength { if len(command) > maxCommandLength {
return ErrorResult(fmt.Sprintf("command too long: %d bytes (max %d)", len(command), maxCommandLength)) return ErrorResult(fmt.Sprintf("command too long: %d bytes (max %d)", len(command), maxCommandLength))
@ -193,6 +195,12 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *To
if strings.TrimSpace(command) == "" { if strings.TrimSpace(command) == "" {
return ErrorResult("command cannot be empty") return ErrorResult("command cannot be empty")
} }
if strings.TrimSpace(command) == ":" {
return ErrorResult("command cannot be a shell no-op placeholder")
}
if guard := predictableLongRunningCommand(command, t.timeout); guard != "" {
return ErrorResult(guard)
}
cwd := t.workingDir cwd := t.workingDir
if wd, ok := args["working_dir"].(string); ok && wd != "" { if wd, ok := args["working_dir"].(string); ok && wd != "" {
@ -462,3 +470,28 @@ func (t *ExecTool) SetAllowPatterns(patterns []string) error {
} }
return nil return nil
} }
func normalizeShellCommand(command string) string {
cmd := strings.TrimSpace(command)
cmd = strings.TrimLeft(cmd, " \t\r\n}")
return strings.TrimSpace(cmd)
}
func predictableLongRunningCommand(command string, timeout time.Duration) string {
fields := strings.Fields(command)
if len(fields) != 2 || fields[0] != "sleep" {
return ""
}
seconds, err := strconv.Atoi(fields[1])
if err != nil || seconds < 0 {
return ""
}
requested := time.Duration(seconds) * time.Second
if requested <= timeout {
return ""
}
return fmt.Sprintf("command denied: requested sleep %ds exceeds timeout budget of %ds", seconds, int(timeout/time.Second))
}

View file

@ -72,7 +72,7 @@ func TestShellTool_Timeout(t *testing.T) {
ctx := t.Context() ctx := t.Context()
args := map[string]interface{}{ args := map[string]interface{}{
"command": "sleep 10", "command": "sh -c 'sleep 10'",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
@ -88,6 +88,55 @@ func TestShellTool_Timeout(t *testing.T) {
} }
} }
func TestShellTool_RejectsNoOpPlaceholder(t *testing.T) {
t.Parallel()
tool := NewExecTool("", false)
result := tool.Execute(t.Context(), map[string]interface{}{
"command": ":",
})
if !result.IsError {
t.Fatal("expected placeholder command to be rejected")
}
if !strings.Contains(result.ForLLM, "no-op placeholder") {
t.Fatalf("expected no-op placeholder error, got %q", result.ForLLM)
}
}
func TestShellTool_NormalizesLeadingGarbage(t *testing.T) {
t.Parallel()
tool := NewExecTool("", false)
result := tool.Execute(t.Context(), map[string]interface{}{
"command": "}\techo normalized-shell-command",
})
if result.IsError {
t.Fatalf("expected normalized command to succeed, got %q", result.ForLLM)
}
if !strings.Contains(result.ForLLM, "normalized-shell-command") {
t.Fatalf("expected normalized command output, got %q", result.ForLLM)
}
}
func TestShellTool_ShortCircuitsSleepBeyondTimeoutBudget(t *testing.T) {
t.Parallel()
tool := NewExecTool("", false)
tool.SetTimeout(8 * time.Second)
result := tool.Execute(t.Context(), map[string]interface{}{
"command": "sleep 120",
})
if !result.IsError {
t.Fatal("expected oversized sleep command to be rejected")
}
if !strings.Contains(result.ForLLM, "exceeds timeout budget") {
t.Fatalf("expected timeout-budget denial, got %q", result.ForLLM)
}
}
// TestShellTool_WorkingDir verifies custom working directory // TestShellTool_WorkingDir verifies custom working directory
func TestShellTool_WorkingDir(t *testing.T) { func TestShellTool_WorkingDir(t *testing.T) {
t.Parallel( t.Parallel(

View file

@ -3,6 +3,7 @@ package tools
import ( import (
"context" "context"
"fmt" "fmt"
"path/filepath"
"strings" "strings"
"sync" "sync"
@ -33,7 +34,7 @@ func (t *SkillSearchTool) getGraph() *skills.SkillGraph {
func (t *SkillSearchTool) Name() string { return "skill_search" } func (t *SkillSearchTool) Name() string { return "skill_search" }
func (t *SkillSearchTool) Description() string { func (t *SkillSearchTool) Description() string {
return "Search available skills by keyword. Returns matching skill names, descriptions, tags, and domains without loading full content. Use this to discover relevant skills before reading them." return "Search available skills by keyword. Returns matching skill names, descriptions, tags, and domains without loading full content. Use this directly for skill discovery instead of tool_search."
} }
func (t *SkillSearchTool) Parameters() map[string]interface{} { func (t *SkillSearchTool) Parameters() map[string]interface{} {
@ -96,7 +97,7 @@ func NewSkillReadTool(loader *skills.SkillsLoader) *SkillReadTool {
func (t *SkillReadTool) Name() string { return "skill_read" } func (t *SkillReadTool) Name() string { return "skill_read" }
func (t *SkillReadTool) Description() string { func (t *SkillReadTool) Description() string {
return "Load the full content of a specific skill by name. Use skill_search first to find relevant skill names, then read the ones you need." return "Load the full content of a specific skill by name. Use skill_search first to find relevant skill names, then call skill_read with {\"name\":\"skill-name\"}."
} }
func (t *SkillReadTool) Parameters() map[string]interface{} { func (t *SkillReadTool) Parameters() map[string]interface{} {
@ -113,11 +114,20 @@ func (t *SkillReadTool) Parameters() map[string]interface{} {
} }
func (t *SkillReadTool) Execute(_ context.Context, args map[string]interface{}) *ToolResult { func (t *SkillReadTool) Execute(_ context.Context, args map[string]interface{}) *ToolResult {
name, _ := args["name"].(string) name := firstNonEmptyString(args, "name", "skill_name", "skill", "path", "query", "title")
if name == "" && len(args) == 1 {
for _, raw := range args {
if s, ok := raw.(string); ok {
name = s
break
}
}
}
if name == "" { if name == "" {
return ErrorResult("name is required") return ErrorResult("name is required")
} }
name = t.normalizeSkillName(name)
content, ok := t.loader.LoadSkill(name) content, ok := t.loader.LoadSkill(name)
if !ok { if !ok {
return ErrorResult(fmt.Sprintf("skill '%s' not found", name)) return ErrorResult(fmt.Sprintf("skill '%s' not found", name))
@ -126,6 +136,78 @@ func (t *SkillReadTool) Execute(_ context.Context, args map[string]interface{})
return SilentResult(fmt.Sprintf("# Skill: %s\n\n%s", name, content)) return SilentResult(fmt.Sprintf("# Skill: %s\n\n%s", name, content))
} }
func (t *SkillReadTool) normalizeSkillName(raw string) string {
name := strings.TrimSpace(raw)
name = strings.Trim(name, "\"'` \t\r\n:{}")
if name == "" {
return name
}
if _, ok := t.loader.LoadSkill(name); ok {
return name
}
cleaned := filepath.ToSlash(name)
cleaned = strings.TrimSuffix(cleaned, "/SKILL.md")
cleaned = strings.TrimSuffix(cleaned, "SKILL.md")
cleaned = strings.Trim(cleaned, "/")
if cleaned != "" {
if _, ok := t.loader.LoadSkill(cleaned); ok {
return cleaned
}
}
parts := strings.Split(cleaned, "/")
for i := len(parts) - 1; i >= 0; i-- {
part := strings.TrimSpace(parts[i])
part = strings.Trim(part, "\"'`")
if part == "" || strings.EqualFold(part, "SKILL.md") {
continue
}
if _, ok := t.loader.LoadSkill(part); ok {
return part
}
}
lower := strings.ToLower(cleaned)
for _, info := range t.loader.ListSkills() {
if strings.Contains(lower, strings.ToLower(info.Name)) {
return info.Name
}
}
var partialMatches []string
for _, info := range t.loader.ListSkills() {
haystack := strings.ToLower(info.Name + " " + info.Description + " " + info.Domain + " " + strings.Join(info.Tags, " "))
if lower != "" && strings.Contains(haystack, lower) {
partialMatches = append(partialMatches, info.Name)
}
}
if len(partialMatches) == 1 {
return partialMatches[0]
}
return name
}
func firstNonEmptyString(args map[string]interface{}, keys ...string) string {
for _, key := range keys {
raw, ok := args[key]
if !ok {
continue
}
s, ok := raw.(string)
if !ok {
continue
}
s = strings.TrimSpace(s)
if s != "" {
return s
}
}
return ""
}
// SkillTraverseTool follows wikilink chains from a skill node. // SkillTraverseTool follows wikilink chains from a skill node.
// This is the third step: after reading a skill, explore its connections. // This is the third step: after reading a skill, explore its connections.
type SkillTraverseTool struct { type SkillTraverseTool struct {

View file

@ -98,6 +98,30 @@ func TestSkillReadTool(t *testing.T) {
assert.Contains(t, result.ForLLM, "position-sizing") assert.Contains(t, result.ForLLM, "position-sizing")
}) })
t.Run("normalizes path-like skill names", func(t *testing.T) {
result := tool.Execute(t.Context(), map[string]interface{}{"name": ".assets/risk-management/SKILL.md"})
assert.False(t, result.IsError)
assert.Contains(t, result.ForLLM, "# Skill: risk-management")
})
t.Run("trims punctuation wrappers around skill names", func(t *testing.T) {
result := tool.Execute(t.Context(), map[string]interface{}{"name": ":\trisk-management:"})
assert.False(t, result.IsError)
assert.Contains(t, result.ForLLM, "# Skill: risk-management")
})
t.Run("accepts alias argument names", func(t *testing.T) {
result := tool.Execute(t.Context(), map[string]interface{}{"skill_name": "risk-management"})
assert.False(t, result.IsError)
assert.Contains(t, result.ForLLM, "# Skill: risk-management")
})
t.Run("falls back from descriptive query when unique", func(t *testing.T) {
result := tool.Execute(t.Context(), map[string]interface{}{"query": "engineering"})
assert.False(t, result.IsError)
assert.Contains(t, result.ForLLM, "# Skill: code-review")
})
t.Run("error on missing skill", func(t *testing.T) { t.Run("error on missing skill", func(t *testing.T) {
result := tool.Execute(t.Context(), map[string]interface{}{"name": "nonexistent"}) result := tool.Execute(t.Context(), map[string]interface{}{"name": "nonexistent"})
assert.True(t, result.IsError) assert.True(t, result.IsError)