diff --git a/.gitignore b/.gitignore index 65ff0af39..649c3cfa0 100644 --- a/.gitignore +++ b/.gitignore @@ -64,3 +64,7 @@ CLAUDE.md # Internal dev setup (not for public repo) /scripts/dev_setup_internal.sh + +# External working trees +/examples/ +/ext/ diff --git a/V0.9_COMPATIBILITY_NOTE.md b/V0.9_COMPATIBILITY_NOTE.md new file mode 100644 index 000000000..0d0018c77 --- /dev/null +++ b/V0.9_COMPATIBILITY_NOTE.md @@ -0,0 +1,132 @@ + + +# V0.9 agentic-runtime Compatibility Note + +本文列出现有用户从 V0.8.x 升级到 V0.9 `agentic-runtime` 时需要关注的 API 和语义变化。未列出的新增能力通常不影响既有 `*schema.Message` 路径。 + +## API 显式变更 + +### ChatModelAgentMiddleware 新增 AfterAgent + +`ChatModelAgentMiddleware` 新增 `AfterAgent` 方法。手写实现该接口的类型需要补充该方法，否则会编译失败。 + +推荐做法： + +- 如果 middleware 不需要特殊收尾逻辑，嵌入 `*adk.BaseChatModelAgentMiddleware`。 +- 如果 middleware 需要在 Agent 成功结束后清理状态、记录事件或补充统计，实现 `AfterAgent(ctx, state)`。 + +影响范围： + +- 仅影响显式实现 `ChatModelAgentMiddleware` 的用户代码。 +- 通过 `BaseChatModelAgentMiddleware` 组合扩展的代码可保持兼容。 + +### summarization.SummarizeMessages 被移除 + +`summarization.SummarizeMessages` 和 `summarization.SummarizeOutput` 不再导出。 + +迁移方式： + +- 构造 summarization middleware 时继续使用 `summarization.New` 或 `summarization.NewTyped`。 +- 需要主动触发同步 summarization 时，使用 `TypedMiddleware.Summarize`。 + +该调整将 summarization 的配置、状态读取和执行逻辑收敛到 middleware 内部，避免独立函数与运行时状态语义分叉。 + +## 需要关注语义变化的能力 + +### Summarization Finalize 后处理语义变化 + +V0.8.x 中，summarization middleware 会先执行默认 summary 后处理，再调用用户配置的 `Finalize`。因此自定义 `Finalize` 收到的 `summary` 已经包含 `PreserveUserMessages` 替换、`TranscriptFilePath` 注入和 summary preamble。 + +V0.9 中，如果设置了 `Config.Finalize`，middleware 会直接把模型生成的 raw summary 传给 `Finalize`，不再自动执行默认后处理。受影响的配置包括： + +- `PreserveUserMessages` +- `TranscriptFilePath` + +迁移方式： + +- 如果希望保留默认后处理，不要设置 `Finalize`，让 middleware 使用默认 finalization 路径。 +- 如果必须自定义 `Finalize`，但仍希望保留默认后处理，先通过 `DefaultFinalizer` 构造默认 finalizer，再在自定义逻辑中显式组合。 +- `DefaultFinalizer` 不会自动读取外层 `Config.PreserveUserMessages` 和 `Config.TranscriptFilePath`；需要通过 `DefaultFinalizerConfig` 显式传入。 +- 使用 `NewFinalizer().PreserveSkills(...).Build()` 的代码需要特别检查：该 finalizer 只负责 preserve skills，不会自动补上 `PreserveUserMessages` 和 `TranscriptFilePath`。 + +### 工具列表修改路径调整 + +`ModelContext.Tools` 不再是推荐的工具列表修改入口。 + +升级建议： + +- 在 `BeforeModelRewriteState` 中修改 `state.ToolInfos`。 +- 如需模型原生 deferred tool search，修改 `state.DeferredToolInfos`。 +- 不建议在 `WrapModel` 中修改工具列表；该修改只影响当前模型调用，后续 middleware、后续 turn 或 checkpoint/resume 不会继承这次修改。 + +### Model Retry 决策语义增强 + +`ModelRetryConfig` 新增 `ShouldRetry`。当 `ShouldRetry` 非空时，`IsRetryAble` 会被忽略。 + +需要注意： + +- 旧的 `IsRetryAble` 仍可用于错误维度的简单重试。 +- 使用 `ShouldRetry` 后，应显式处理成功输出但业务不接受的场景。 +- Interrupt 和 `ErrStreamCanceled` 不作为普通 retry error 处理。 + +### Cancel 错误语义 + +V0.9 引入主动取消语义后，应用需要区分主动取消、普通错误和业务 interrupt。 + +升级建议： + +- 上层应区分 `CancelError`、普通 error 和业务 interrupt。 +- 如果应用主动接入 `WithCancel`，不要把 `CancelError` 当作普通业务失败处理。 + +### AgenticMessage 迁移需要理解新的消息结构 + +`TypedChatModelAgent[*schema.AgenticMessage]` 是面向模型原生 Agentic 协议的新路径。迁移到该路径不只是把泛型参数从 `*schema.Message` 改成 `*schema.AgenticMessage`，还需要按 `AgenticMessage` 的 content block 结构处理消息内容。 + +需要注意： + +- AgenticMessage 路径使用 `AgenticModel` 与 `AgenticToolsNode` 处理工具调用。 +- 工具调用和工具结果通过 `AgenticMessage` content block 表达，尤其需要正确处理 tool call / tool result content block。 +- Agent transfer 能力不适用于 AgenticMessage 路径。 +- 既有应用如果不需要模型原生 Agentic 协议，建议继续使用默认 `*schema.Message` 路径；只有在明确要接入 `AgenticModel` 协议时再迁移。 + +### 模型适配器需要识别新增 option + +V0.9 引入 `AgenticModel` 后，模型适配器需要更严格地处理 call-time options。`AgenticModel` 是 `BaseModel[*schema.AgenticMessage]` 的别名，不再提供类似 `ToolCallingChatModel.WithTools` 的增强接口；工具绑定统一通过 `model.WithTools` 作为 `model.Option` 传入。 + +需要注意： + +- 所有支持 AgenticMessage 的模型适配器都应读取 `Options.Tools`，并将其映射到 provider 的 tool calling 协议。 +- `AgenticModel` 不应要求用户先调用某个 `WithTools` 方法得到“带工具的模型实例”；ADK 会在每次模型调用时通过 `model.WithTools` 传递当前工具列表。 +- 如果适配器只从自身 config 读取工具，而忽略 `model.WithTools`，在 ChatModelAgent / AgenticToolsNode 路径下会出现模型看不到工具或工具列表不随运行态变化的问题。 + +V0.9 还在 `model.Options` 中新增： + +- `DeferredTools` +- `ToolSearchTool` +- `AgenticToolChoice` + +现有模型适配器忽略这些 option 通常不会导致编译失败，但会导致 deferred tool search、模型原生 tool search 或 agentic tool choice 不生效。适配器维护者应按目标 provider 的协议补齐转换逻辑。 + +### ToolInfo 序列化形态变化 + +`ToolInfo` 增加显式 JSON/Gob 编解码，以保留 `ParamsOneOf`。 + +影响： + +- `ToolInfo` 进入了 `ChatModelAgentState.ToolInfos` / `DeferredToolInfos`，因此可能随 Agent state 一起进入 checkpoint。 +- 显式 JSON/Gob 编解码用于保证 `ParamsOneOf` 在 checkpoint、deep copy 和恢复过程中不会丢失。 +- 如果外部系统直接依赖旧版 `ToolInfo` JSON 形态，需要重新确认序列化兼容性。 diff --git a/V0.9_RELEASE_FINDINGS.md b/V0.9_RELEASE_FINDINGS.md new file mode 100644 index 000000000..e58c2aaa3 --- /dev/null +++ b/V0.9_RELEASE_FINDINGS.md @@ -0,0 +1,90 @@ + + +# v0.9 Release Findings + +## Comparison Scope + +- Compared `alpha/09` against `main` using `main...alpha/09`. +- `main` is the merge base of `alpha/09`. +- Branch heads observed during analysis: + - `main`: `5e1305506c4fa89ef5d786035a947258e29a7593` + - `alpha/09`: `c39433511896d6a12e379a7958c6e5d489560b5a` +- Second validation pass confirmed `main == origin/main`, `alpha/09 == origin/alpha/09`, and `main` is the merge base. +- Diff size: `136 files changed`, `49,967 insertions`, `2,790 deletions`. +- Changed surface is concentrated in `adk`, `schema`, `components/model`, `components/prompt`, `compose`, and callback helpers. + +## Primary Features + +| Area | v0.9 feature | Direct diff validation | +| --- | --- | --- | +| Agentic message model | Adds `schema.AgenticMessage`, content-block based message schema, provider extensions, streaming metadata, MCP/server/function tool blocks, and concat support. | `A schema/agentic_message.go`; concat registration in `schema/message.go`. | +| Generic model abstraction | Introduces `model.BaseModel[M]`; keeps `BaseChatModel` as `BaseModel[*schema.Message]`; adds `AgenticModel`. | `M components/model/interface.go`. | +| Typed ADK | Adds typed agents, typed events, typed runner, typed `ChatModelAgent`, and typed message variants while preserving default `*schema.Message` aliases. | `M adk/interface.go`, `M adk/chatmodel.go`, `M adk/runner.go`. | +| Agentic ChatModelAgent path | `TypedChatModelAgent[*schema.AgenticMessage]` supports a single-shot agentic model path where tool calling is handled inside the model/message protocol. | `M adk/chatmodel.go`; `TypedChatModelAgent` and agentic ReAct path are added in the diff. | +| Cancellation | Adds `WithCancel`, `CancelMode`, safe-point cancellation, recursive cancellation, timeout escalation, `CancelHandle`, and `CancelError` with resumable interrupt contexts. | `A adk/cancel.go`; cancel integration hunks in `adk/chatmodel.go`, `adk/flow.go`, and `adk/wrappers.go`. | +| TurnLoop | Adds a push-based `TurnLoop` runtime with `Push`, non-blocking `Stop`, idle-stop, checkpoint/resume integration, and preempt handling. | `A adk/turn_loop.go`, `A adk/turn_buffer.go`. | +| Model retry | Upgrades retry from error-only retryability to `ShouldRetry(ctx, RetryContext) -> RetryDecision`, allowing output inspection, input rewrite, option rewrite, backoff override, and reject reason. | `M adk/retry_chatmodel.go`. | +| Model failover | Adds ChatModel failover with `ModelFailoverConfig`, `FailoverContext`, last-success model preference, and callback-aware proxying. | `A adk/failover_chatmodel.go`; config wiring in `adk/chatmodel.go`. | +| Tool search | Adds dynamic tool search middleware with both client-side search and model-native deferred tool search via `DeferredToolInfos`, `WithDeferredTools`, and `WithToolSearchTool`. | `M adk/middlewares/dynamictool/toolsearch/toolsearch.go`, `M components/model/option.go`. | +| Middleware modernization | Generifies summarization, reduction, skill, filesystem, plan-task, patch-tool-calls and adds `AfterAgent`; state now carries `ToolInfos` and `DeferredToolInfos` as the recommended mutable model-call surface. | Diff hunks in `adk/handler.go`, `adk/middlewares/*`, and `adk/prebuilt/deep/deep.go`. | +| Summarization API | Adds `TypedMiddleware.Summarize` and typed finalizer/customized-action paths; removes the old standalone `SummarizeMessages` / `SummarizeOutput` API in favor of middleware-owned summarization. | `M adk/middlewares/summarization/summarization.go`, `M customized_action.go`, `M finalizer_builder.go`. | +| Compose/tooling | Adds `AgenticToolsNode` and tool name/argument aliases for `ToolsNode`. | `A compose/agentic_tools_node.go`, `M compose/tool_node.go`. | +| Prompt/callback support | Adds agentic prompt templates and callback types for agentic prompt/model/tools/agent components. | `A components/prompt/agentic_chat_template.go`, `A components/*/agentic_callback_extra.go`, `M utils/callbacks/template.go`. | +| Filesystem | Adds enhanced multimodal read support and PDF page validation. | `M adk/filesystem/backend.go`, `M adk/middlewares/filesystem/filesystem.go`. | +| Agents.md | Adds `agentsmd` middleware for automatically loading and injecting `AGENTS.md`-style instructions. | `A adk/middlewares/agentsmd/agentsmd.go`, `A loader.go`. | + +## Compatibility Notes + +| Impact | Note | +| --- | --- | +| Source break for custom middleware implementers | `ChatModelAgentMiddleware` now includes `AfterAgent`. Any user-defined type that manually implements the interface must add this method or embed `BaseChatModelAgentMiddleware`. | +| Middleware tool mutation semantics | `ModelContext.Tools` is now deprecated as a mutation surface; tool list changes should happen through `state.ToolInfos` / `state.DeferredToolInfos` in `BeforeModelRewriteState`. Mutating tools in `WrapModel` only affects one model call and is explicitly discouraged. | +| Summarization standalone API removal | `summarization.SummarizeMessages` and `summarization.SummarizeOutput` are no longer exported. Use `New` / `NewTyped` to construct middleware, or call `TypedMiddleware.Summarize` when direct summarization is needed. | +| Retry behavior change | If `ShouldRetry` is set, `IsRetryAble` is ignored. In streaming mode, the full stream is consumed before the retry decision is made, although events are still emitted in real time. | +| Retry cancellation semantics | Retry now treats interrupts and `ErrStreamCanceled` as non-retryable and uses context-aware backoff rather than unconditional sleep. Users relying on retrying interrupt/cancel errors should adjust policy. | +| Cancellation error semantics | During active cancel, business interrupts are absorbed into `CancelError`; the checkpoint preserves interrupt contexts and business interrupt can re-fire on resume. Consumers should handle `CancelError` separately from ordinary business interrupts. | +| TurnLoop stop semantics | `TurnLoop.Stop` is non-blocking; use `Wait` for terminal state. Cancel-related stop options degrade to "finish current turn then exit" if the running agent does not support `WithCancel`. `UntilIdleFor` silently drops cancel options in the same call. | +| Agentic path limitations | `TypedChatModelAgent[*schema.AgenticMessage]` is not feature-equivalent to `*schema.Message`: it uses a single-shot path, does not support agent transfer, and cancel monitoring/retry on model streams are not yet wired. | +| Model adapters must honor new options | Native tool search requires model implementations to read `Options.DeferredTools`, `Options.ToolSearchTool`, and `Options.AgenticToolChoice`. Existing adapters that ignore unknown common options will compile but will not support the new behavior. | +| Serialization shape change | `ToolInfo` now has explicit JSON/Gob encoding that preserves `ParamsOneOf`. This fixes checkpoint/deep-copy loss, but external systems depending on the previous raw JSON shape should re-check serialized payloads. | +| Filesystem page validation | Multimodal read validates PDF `pages` and rejects ranges over 20 pages per request. Users passing arbitrary page ranges should handle validation errors. | +| Transfer/workflow/supervisor positioning | Agent transfer, workflow agents, and supervisor are not removed, but many APIs now carry `NOT RECOMMENDED` guidance in favor of `ChatModelAgent` + `AgentTool` or `DeepAgent`. This is a semantic/product-direction compatibility note, not a signature break. | + +## Likely Non-Breaking Alias Changes + +- `BaseChatModel` becomes an alias of `BaseModel[*schema.Message]`; existing implementations with `Generate(ctx, []*schema.Message, ...)` and `Stream(ctx, []*schema.Message, ...)` should still satisfy it. +- `Agent`, `AgentInput`, `AgentEvent`, `AgentOutput`, `ChatModelAgent`, `ChatModelAgentConfig`, `ChatModelAgentState`, `ModelContext`, and several middleware config types are preserved as `*schema.Message` aliases over typed forms. +- `ToolOutputPart`, `ToolResult`, and related tool-result types moved from `schema/message.go` to `schema/tool.go`, but remain in package `schema`, so import paths and qualified names are unchanged. + +## Validation Results + +Completed checks: + +- Direct branch-ref validation: + - Verified `main == origin/main`, `alpha/09 == origin/alpha/09`, and the merge base is `main`. + - Rechecked each retained feature row with `git diff main...alpha/09` file status or hunks. + - Removed raw AST API-diff counts from the release findings because the script over-reported generic alias refactors as removals. +- Representative downstream compatibility compile check: + - `GOWORK=off go test .` passed in a temporary external module using `replace github.com/cloudwego/eino => ..`. + - Verified that existing `BaseChatModel` implementations still compile against the `BaseModel[*schema.Message]` alias. + - Verified that `ChatModelAgentConfig`, `summarization.Config`, `reduction.Config`, `skill.Config`, `ToolResult`, and new model options are usable from downstream code. + - Verified that embedding `*adk.BaseChatModelAgentMiddleware` remains the safe compatibility path for middleware implementations. +- Negative compile check for old custom middleware: + - `GOWORK=off go test -tags=oldmiddleware .` fails as expected with: `oldStyleMiddleware does not implement adk.TypedChatModelAgentMiddleware[*schema.Message] (missing method AfterAgent)`. + - This confirms the `AfterAgent` source compatibility note for users who manually implement `ChatModelAgentMiddleware` without embedding the base middleware. +- Targeted package tests: + - `go test ./adk ./adk/middlewares/summarization ./adk/middlewares/reduction ./adk/middlewares/skill ./adk/middlewares/dynamictool/toolsearch ./components/model ./components/prompt ./compose ./schema` passed. diff --git a/V0.9_RELEASE_NOTE.md b/V0.9_RELEASE_NOTE.md new file mode 100644 index 000000000..dc50213ef --- /dev/null +++ b/V0.9_RELEASE_NOTE.md @@ -0,0 +1,70 @@ + + +# V0.9 agentic-runtime Release Note + +V0.9 的版本主题是 `agentic-runtime`。该版本主要围绕 ADK 的消息协议、Agent 运行控制和多轮运行时能力展开，在保留 `*schema.Message` 默认路径的同时，引入 `AgenticMessage` 及配套泛型抽象，为更丰富的模型原生 Agent 协议、服务端工具调用、运行中断与恢复打下基础。 + +## 1. AgenticMessage 与 ADK 支持 + +V0.9 新增 `schema.AgenticMessage`，用于表达比传统 `schema.Message` 更完整的 Agentic 消息结构。 + +- `AgenticMessage` 采用 content block 模型，支持文本、推理内容、工具调用、工具结果、服务端工具、MCP 工具和多模态内容等结构化片段。 +- `[]ContentBlock` 能更完整地保留不同模型协议响应中的 block 时序；新增 block 类型也更适配 OpenAI Responses API、Claude、Gemini 等协议中的 tool use、reasoning、streaming metadata 等结构。 +- `components/model` 新增 `AgenticModel` 组件，用于接入以 `AgenticMessage` 为输入输出的模型实现。 +- ADK 对 `AgenticMessage` 路径提供 typed agent、typed event、typed runner 和 typed `ChatModelAgent` 支持，使 AgenticModel 能进入 ADK 的 Agent 生命周期。 + +## 2. ChatModelAgent 能力扩展 + +V0.9 对 `ChatModelAgent` 的运行控制、模型调用可靠性和 middleware 扩展点进行了系统增强。 + +### Cancel + +- 新增 Agent Cancel 能力，用于从外部主动终止正在运行的 Agent。 +- 支持安全点取消、递归取消、取消超时升级，以及取消过程中的 checkpoint 持久化。 +- 取消期间发生的 interrupt 会统一进入取消语义，调用方可以通过 `CancelError` 区分主动取消与普通业务失败。 + +### Model Retry + +- Retry 从简单的 error retry 扩展为 `ShouldRetry(ctx, RetryContext) -> RetryDecision`。 +- Retry 决策可以读取模型输出、拒绝不满足条件的输出、修改下一次输入、追加模型 option，并覆盖 backoff。 + +### Model Failover + +- 新增 Model Failover 能力，用于在模型调用失败后切换到备用模型。 +- Failover 决策可以读取失败 attempt 的输出、错误、原始输入和 attempt 序号，并选择下一次使用的模型。 +- 支持为备用模型改写输入；也支持优先复用上一次调用成功的模型，降低每次从固定主模型开始试错的成本。 + +### Middleware 增强 + +- `ChatModelAgentMiddleware` 新增 `AfterAgent`，用于在 Agent 成功结束后执行收尾逻辑。 +- Summarization、reduction、skill、filesystem、plan-task、patch-tool-calls 等 middleware 完成泛型化，支持 `AgenticMessage` 路径。 +- Summarization middleware 新增 `TypedMiddleware.Summarize`，同步 summarization 能力从独立函数转为 middleware 内聚能力。 +- Filesystem middleware 增强多模态读取能力，并增加 PDF pages 校验。 +- 新增 `agentsmd` middleware，用于加载和注入 `AGENTS.md` 风格的项目指令。 +- `ChatModelAgentState` 增加 `ToolInfos` 和 `DeferredToolInfos`，作为 middleware 调整模型可见工具集合的主路径。 +- `ToolInfos` 表示当前模型调用直接可见的工具；`DeferredToolInfos` 表示可由模型通过工具搜索机制按需发现的候选工具。 +- Tool search middleware 支持三类工具加载方式：使用模型侧原生 tool search 能力从 deferred tools 中按需加载；按模型协议要求提供固定 schema 的 `ToolSearchTool`，由模型通过该入口搜索 deferred tools；不依赖模型侧协议，使用 Eino 提供的自定义 `tool_search` tool 检索工具，并把命中的工具追加到常规 `ToolInfos`。 +- Compose 新增 `AgenticToolsNode`，`ToolsNode` 增加 tool name 和 argument alias 支持。 + +## 3. TurnLoop + +V0.9 新增 `TurnLoop`，用于把一次性的 Agent run 提升为可持续运行、可被外部驱动的 turn 级运行时。 + +- 面向多轮运行：`TurnLoop` 持续接收外部输入，每个 turn 独立规划输入、构造 Agent、消费事件，适合长期在线的交互式 Agent。 +- 支持输入合并：`GenInput` 在 turn 边界决定本轮消费哪些输入、哪些继续等待，应用可以实现批处理、去重、合并用户连续输入等策略。 +- 支持抢占：带 preempt option 的 `Push` 会原子地写入新输入并请求取消当前 turn，使高优先级输入可以打断正在运行的 Agent。 +- 支持声明式 checkpoint/resume：恢复时，应用不需要自行还原输入队列；`TurnLoop` 会区分被中断的输入、尚未处理的输入和恢复后新到达的输入，应用只需声明这些输入如何重新进入后续 turn。 diff --git a/adk/agent_tool.go b/adk/agent_tool.go index 3f120c238..0d2c40565 100644 --- a/adk/agent_tool.go +++ b/adk/agent_tool.go @@ -19,10 +19,12 @@ package adk import ( "context" + "encoding/json" "errors" "fmt" "github.com/bytedance/sonic" + "github.com/google/uuid" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/compose" @@ -151,6 +153,15 @@ func (at *typedAgentTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) }, nil } +// agentToolInterruptState is the JSON-encoded state captured when an AgentTool +// invocation is interrupted. It wraps the bridge checkpoint bytes alongside +// the synthetic child session ID so resume preserves SessionID-based event +// filtering across interrupt/resume. +type agentToolInterruptState struct { + ChildSessionID string `json:"child_session_id"` + BridgeCheckpoint []byte `json:"bridge_checkpoint"` +} + func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { if cancelCtx := getCancelContext(ctx); cancelCtx != nil { cancelCtx.markAgentToolDescendant() @@ -161,7 +172,35 @@ func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON s var iter *AsyncIterator[*TypedAgentEvent[M]] var err error - wasInterrupted, hasState, state := tool.GetInterruptState[[]byte](ctx) + wasInterrupted, hasState, rawState := tool.GetInterruptState[[]byte](ctx) + + var childSessionID string + var bridgeCheckpoint []byte + + if !wasInterrupted { + // First invocation — generate a globally-unique child session ID. + // Synthetic UUID avoids collisions with model-assigned tool call IDs + // (which may be reused across turns) and with user-assigned session IDs. + childSessionID = "agent_tool:" + uuid.NewString() + } else if !hasState { + return "", fmt.Errorf("agent tool '%s' interrupt has happened, but cannot find interrupt state", at.agent.Name(ctx)) + } else { + // Resume — try the JSON envelope (introduced when SessionID-based event + // filtering landed). If the envelope does not parse or carries no bridge + // checkpoint, the rawState is from a pre-envelope version: treat the + // raw bytes as the bridge checkpoint and synthesize a fresh + // childSessionID. Pre-envelope checkpoints predate session persistence, + // so the synthesized ID has no parent-session filter to coordinate with. + var wrapped agentToolInterruptState + if json.Unmarshal(rawState, &wrapped) == nil && len(wrapped.BridgeCheckpoint) > 0 { + childSessionID = wrapped.ChildSessionID + bridgeCheckpoint = wrapped.BridgeCheckpoint + } else { + childSessionID = "agent_tool:" + uuid.NewString() + bridgeCheckpoint = rawState + } + } + if !wasInterrupted { ms = newBridgeStore() @@ -169,11 +208,6 @@ func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON s if at.fullChatHistoryAsInput { var zero M if _, ok := any(zero).(*schema.Message); !ok { - // fullChatHistoryAsInput is only supported for *schema.Message agents and will not - // be extended to *schema.AgenticMessage. The chat history format and role semantics - // differ fundamentally between Message and AgenticMessage, and the history rewriting - // logic (role attribution, system message filtering, transfer messages) is specific - // to the Message model. return "", fmt.Errorf("fullChatHistoryAsInput is only supported for *schema.Message agents") } msgInput, histErr := getReactChatHistory(ctx, at.agent.Name(ctx)) @@ -197,11 +231,7 @@ func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON s iter = runner.Run(ctx, input, append(extractAndDeriveAgentToolCancelCtx(ctx, at.agent.Name(ctx), opts), WithCheckPointID(bridgeCheckpointID), withSharedParentSession())...) } else { - if !hasState { - return "", fmt.Errorf("agent tool '%s' interrupt has happened, but cannot find interrupt state", at.agent.Name(ctx)) - } - - ms = newResumeBridgeStore(bridgeCheckpointID, state) + ms = newResumeBridgeStore(bridgeCheckpointID, bridgeCheckpoint) agentOpts := extractAndDeriveAgentToolCancelCtx(ctx, at.agent.Name(ctx), opts) agentOpts = append(agentOpts, withSharedParentSession()) @@ -239,6 +269,10 @@ func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON s rp = append(rp, event.RunPath...) event.RunPath = rp } + // Tag forwarded events with the child session ID so live consumers + // can distinguish child timeline events and the parent's persistence + // loop can skip them. + stampAgentToolSessionEvent(event, childSessionID) tmp := copyTypedAgentEvent(event) gen.Send(event) event = tmp @@ -257,7 +291,17 @@ func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON s return "", fmt.Errorf("interrupt has happened, but cannot find interrupt info") } - return "", tool.CompositeInterrupt(ctx, "agent tool interrupt", data, + // Wrap bridge checkpoint with childSessionID so resume can recover it. + wrapped := agentToolInterruptState{ + ChildSessionID: childSessionID, + BridgeCheckpoint: data, + } + wrappedBytes, mErr := json.Marshal(wrapped) + if mErr != nil { + return "", fmt.Errorf("agent_tool: failed to encode interrupt state: %w", mErr) + } + + return "", tool.CompositeInterrupt(ctx, "agent tool interrupt", wrappedBytes, lastEvent.Action.internalInterrupted) } @@ -408,6 +452,41 @@ func newTypedUserMessages[M MessageType](text string) []M { } } +func stampAgentToolSessionEvent[M MessageType](event *TypedAgentEvent[M], childSessionID string) { + if event == nil || childSessionID == "" { + return + } + if event.SessionEventVariant == nil && event.Output != nil && event.Output.MessageOutput != nil { + ts := newEventTimestamp() + if event.Output.MessageOutput.IsStreaming { + event.SessionEventVariant = &SessionEventVariant[M]{ + MessageStreamRef: &MessageStreamRef{ + Timestamp: ts, + Kind: SessionEventMessage, + }, + } + } else if !isNilMessage(event.Output.MessageOutput.Message) { + event.SessionEventVariant = &SessionEventVariant[M]{ + Event: &SessionEvent[M]{ + Timestamp: ts, + Kind: SessionEventMessage, + Message: event.Output.MessageOutput.Message, + }, + } + } + } + if event.SessionEventVariant != nil { + event.SessionEventVariant.SessionID = childSessionID + } +} + +// newTypedInvokableAgentToolRunner creates a runner for the inner agent without +// SessionEventStore. The child's events are forwarded to the parent's live stream +// (tagged with childSessionID on SessionEventVariant) and filtered out of the parent's persistence. +// The child's durability relies solely on the bridge checkpoint stored inside +// agentToolInterruptState — there is no independent child session log. +// This may change in the future if AgentTool needs cross-turn context +// continuation or audit-level event logging for the child session. func newTypedInvokableAgentToolRunner[M MessageType](agent TypedAgent[M], store compose.CheckPointStore, enableStreaming bool) *TypedRunner[M] { return &TypedRunner[M]{ a: agent, diff --git a/adk/agent_tool_test.go b/adk/agent_tool_test.go index 785ad995a..16e768811 100644 --- a/adk/agent_tool_test.go +++ b/adk/agent_tool_test.go @@ -937,6 +937,44 @@ func TestAgentTool_InvokableRun_StreamingVariant(t *testing.T) { } } +func TestStampAgentToolSessionEvent(t *testing.T) { + msg := schema.AssistantMessage("child", nil) + event := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: msg, Role: schema.Assistant}, + }, + } + + stampAgentToolSessionEvent(event, "agent_tool:child") + + require.NotNil(t, event.SessionEventVariant.Event) + assert.Equal(t, "agent_tool:child", event.SessionEventVariant.SessionID) + assert.Empty(t, event.SessionEventVariant.Event.EventID) + assert.False(t, event.SessionEventVariant.Event.Timestamp.IsZero()) + assert.Equal(t, SessionEventMessage, event.SessionEventVariant.Event.Kind) +} + +func TestStampAgentToolSessionEvent_Streaming(t *testing.T) { + event := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: schema.StreamReaderFromArray([]Message{schema.AssistantMessage("child", nil)}), + Role: schema.Assistant, + }, + }, + } + + stampAgentToolSessionEvent(event, "agent_tool:child") + + ref := event.SessionEventVariant.MessageStreamRef + require.NotNil(t, ref) + assert.Equal(t, "agent_tool:child", event.SessionEventVariant.SessionID) + assert.Empty(t, ref.EventID) + assert.False(t, ref.Timestamp.IsZero()) + assert.Equal(t, SessionEventMessage, ref.Kind) +} + func TestSequentialWorkflow_WithChatModelAgentTool_NestedRunPathAndSessions(t *testing.T) { ctx := context.Background() diff --git a/adk/backgroundtask/id.go b/adk/backgroundtask/id.go new file mode 100644 index 000000000..7bcd21fb9 --- /dev/null +++ b/adk/backgroundtask/id.go @@ -0,0 +1,109 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package backgroundtask + +import ( + "math/rand" + "time" +) + +// Task-id layout: a positive int64 (63 usable bits) packed as +// +// [ 41 bits ms timestamp ][ 12 bits sequence ][ 10 bits random ] +// +// Uniqueness within a process is guaranteed by (timestamp, sequence): the +// sequence resets each millisecond and increments for every id minted within the +// same millisecond, all under the Manager lock. If more than 2^12 ids are minted +// in a single millisecond the generator spins to the next millisecond rather than +// wrapping the sequence, so (timestamp, sequence) never repeats. The random low +// bits only make ids look unordered/unpredictable; they are not relied upon for +// uniqueness. +// +// 41 bits of milliseconds covers ~69 years; 12 bits allows 4096 ids per +// millisecond before the generator advances to the next millisecond. +const ( + idSeqBits = 12 + idRandomBits = 10 + idSeqLimit = 1 << idSeqBits + idRandomMask = (1 << idRandomBits) - 1 +) + +// nextRawID packs the next task id integer. Must be called with m.mu held, as it +// reads and advances m.seq / m.lastMs. +func (m *Manager) nextRawID() int64 { + ms := time.Now().UnixMilli() + switch { + case ms > m.lastMs: + m.lastMs = ms + m.seq = 0 + default: + // Same millisecond (or a backward clock step): keep the id monotonic by + // staying on lastMs and advancing the sequence. On sequence overflow, move + // to the next millisecond so (timestamp, sequence) stays unique. + ms = m.lastMs + m.seq++ + if m.seq >= idSeqLimit { + ms = m.waitNextMs(m.lastMs) + m.lastMs = ms + m.seq = 0 + } + } + + //nolint:gosec // non-cryptographic: random bits only diffuse the id's look. + r := int64(rand.Intn(idRandomMask + 1)) + return (ms << (idSeqBits + idRandomBits)) | (m.seq << idRandomBits) | r +} + +// waitNextMs busy-waits until the wall clock advances past prevMs. Reached only +// when more than 2^12 ids are minted within one millisecond. +func (m *Manager) waitNextMs(prevMs int64) int64 { + ms := time.Now().UnixMilli() + for ms <= prevMs { + ms = time.Now().UnixMilli() + } + return ms +} + +// base62 encodes a non-negative int64 using [0-9A-Za-z]. It is the compact, +// URL-safe textual form of a task id's integer. +func base62(n int64) string { + const alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + if n == 0 { + return "0" + } + var buf [11]byte // ceil(63 / log2(62)) = 11 + i := len(buf) + for n > 0 { + i-- + buf[i] = alphabet[n%62] + n /= 62 + } + return string(buf[i:]) +} + +// defaultTaskIDPrefix is used when a task has no Type tag. +const defaultTaskIDPrefix = "task" + +// taskIDPrefix returns the id prefix for a task type, falling back to a generic +// prefix when the type is empty. The type tag (e.g. "bash", "subagent") makes ids +// self-describing: "bash_3Fa9...". +func taskIDPrefix(taskType string) string { + if taskType == "" { + return defaultTaskIDPrefix + } + return taskType +} diff --git a/adk/backgroundtask/id_test.go b/adk/backgroundtask/id_test.go new file mode 100644 index 000000000..dfe3ff8aa --- /dev/null +++ b/adk/backgroundtask/id_test.go @@ -0,0 +1,160 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package backgroundtask + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBase62(t *testing.T) { + assert.Equal(t, "0", base62(0)) + assert.Equal(t, "A", base62(10)) + assert.Equal(t, "10", base62(62)) + // Round-trippable shape: only alphabet chars, non-empty. + const alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + for _, n := range []int64{1, 61, 100, 1 << 40, (1 << 63) - 1} { + s := base62(n) + assert.NotEmpty(t, s) + for _, c := range s { + assert.True(t, strings.ContainsRune(alphabet, c), "char %q not in alphabet", c) + } + } +} + +func TestTaskIDPrefix(t *testing.T) { + assert.Equal(t, "bash", taskIDPrefix("bash")) + assert.Equal(t, "subagent", taskIDPrefix("subagent")) + assert.Equal(t, defaultTaskIDPrefix, taskIDPrefix("")) +} + +// IDs minted in a tight loop within one process must never collide, and must +// carry the task-type prefix. +func TestCreateTask_IDsUniqueAndPrefixed(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + const n = 20000 + seen := make(map[string]struct{}, n) + for i := 0; i < n; i++ { + id, err := m.createTask(context.Background(), &RunInput{Type: "bash", Description: "x"}) + if err != nil { + t.Fatalf("createTask: %v", err) + } + assert.True(t, strings.HasPrefix(id, "bash_"), "id %q missing type prefix", id) + if _, dup := seen[id]; dup { + t.Fatalf("duplicate id generated: %q", id) + } + seen[id] = struct{}{} + } + assert.Len(t, seen, n) +} + +// An empty task type falls back to the generic prefix. +func TestCreateTask_EmptyTypePrefix(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + id, err := m.createTask(context.Background(), &RunInput{Description: "x"}) + assert.NoError(t, err) + assert.True(t, strings.HasPrefix(id, defaultTaskIDPrefix+"_"), "id %q", id) +} + +type taskIDContextKey struct{} + +func TestManager_IDGenOverridesDefaultID(t *testing.T) { + const wantID = "short_000001" + ctx := context.WithValue(context.Background(), taskIDContextKey{}, "trace-1") + called := false + m := New(context.Background(), &Config{ + IDGen: func(ctx context.Context, input *RunInput) (string, error) { + called = true + assert.Equal(t, "bash", input.Type) + assert.Equal(t, "call_1", input.ToolUseID) + assert.Equal(t, "trace-1", ctx.Value(taskIDContextKey{})) + return wantID, nil + }, + }) + defer closeWithTimeout(m) + + result, err := m.Run(ctx, &RunInput{ + Description: "x", + Type: "bash", + ToolUseID: "call_1", + }, workReturning("ok", nil)) + require.NoError(t, err) + assert.True(t, called) + assert.Equal(t, wantID, result.ID) + + task, ok := m.Get(wantID) + require.True(t, ok) + assert.Equal(t, wantID, task.ID) + assert.Equal(t, "bash", task.Type) +} + +func TestCreateTask_IDGenEmptyIDFails(t *testing.T) { + m := New(context.Background(), &Config{ + IDGen: func(context.Context, *RunInput) (string, error) { + return "", nil + }, + }) + defer closeWithTimeout(m) + + _, err := m.createTask(context.Background(), &RunInput{Description: "x"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty id") + assert.Empty(t, m.List()) +} + +func TestCreateTask_IDGenDuplicateIDFails(t *testing.T) { + m := New(context.Background(), &Config{ + IDGen: func(context.Context, *RunInput) (string, error) { + return "fixed", nil + }, + }) + defer closeWithTimeout(m) + + id, err := m.createTask(context.Background(), &RunInput{Description: "first"}) + require.NoError(t, err) + assert.Equal(t, "fixed", id) + + _, err = m.createTask(context.Background(), &RunInput{Description: "second"}) + require.Error(t, err) + assert.Contains(t, err.Error(), `task id "fixed" already exists`) + assert.Len(t, m.List(), 1) +} + +func TestManager_IDGenErrorFailsRun(t *testing.T) { + wantErr := errors.New("allocate id") + m := New(context.Background(), &Config{ + IDGen: func(context.Context, *RunInput) (string, error) { + return "", wantErr + }, + }) + defer closeWithTimeout(m) + + _, err := m.Run(context.Background(), &RunInput{Description: "x"}, workReturning("ok", nil)) + require.Error(t, err) + assert.ErrorIs(t, err, wantErr) + assert.Contains(t, err.Error(), "task id generator") + assert.Empty(t, m.List()) +} diff --git a/adk/backgroundtask/manager.go b/adk/backgroundtask/manager.go new file mode 100644 index 000000000..153b88168 --- /dev/null +++ b/adk/backgroundtask/manager.go @@ -0,0 +1,1282 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package backgroundtask provides a shared lifecycle registry for long-running +// executions (sub-agents, shell commands, ...) that may outlive the tool call +// that launched them. +// +// The central type is Manager: a non-generic, in-memory registry that tracks +// foreground/background/auto-background runs and exposes them via Get/List/ +// Cancel/Wait/Close. Manager is deliberately non-generic so a single instance +// can be shared across heterogeneous domains (e.g. agent runs and shell runs) +// under one unified task-ID space. +// +// What a run actually does is supplied per-call as a WorkFunc passed to Run, which +// keeps the engine agnostic to what it runs (agents, shell, ...); adapters that +// produce WorkFunc values live in the consuming packages (subagent, filesystem). +// A task may carry an output-file path (RunInput.OutputFile) that the launcher +// writes to; the Manager only records and surfaces the path, never the file. +// +// Manager tracks lifecycle only. Streaming a specific run's events in real time +// is intentionally out of scope here: the launcher (a domain adapter) already +// knows the concrete event type, so live streaming, if needed, belongs at that +// typed layer rather than behind a type-erased registry-wide channel. +package backgroundtask + +import ( + "context" + "fmt" + "io" + "runtime/debug" + "strings" + "sync" + "time" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +// Status represents the lifecycle status of a task. +type Status string + +const ( + // StatusRunning indicates the task is currently executing. + StatusRunning Status = "running" + // StatusCompleted indicates the task finished successfully. + StatusCompleted Status = "completed" + // StatusFailed indicates the task terminated with an error. + StatusFailed Status = "failed" + // StatusCanceled indicates the task was stopped by an external request + // (Cancel / Close) via context cancellation. + StatusCanceled Status = "canceled" +) + +// Task represents a single managed execution record. +type Task struct { + // ID is the unique identifier for this task, generated by Manager. + ID string + // Type is a caller-supplied tag (e.g. "bash", "subagent") identifying what kind + // of work this task runs. The Manager does not interpret it; it lets the + // ShouldAutoBackground hook and observers distinguish domains without parsing + // Description. Empty if the launcher did not set it. + Type string + // ToolUseID is the id of the tool call that launched this task, when known. + // It lets a host correlate a task event back to the originating tool call. + // Empty if the launcher did not set it. + ToolUseID string + // Description is a human-readable summary of what the task does. + Description string + // Status is the current lifecycle status. + Status Status + // Result contains the task output, set when Status is StatusCompleted. + Result string + // OutputFile is the path the task's output is written to, when the launcher + // supplied one via RunInput.OutputFile. Empty otherwise. The Manager only + // records and surfaces this path (in notices and Get/List); it never writes + // the file — the launcher's work owns writing, so it may contain interim + // output while the task is still running. The file outlives the in-memory + // record, so it remains readable after the Manager is rebuilt. + OutputFile string + // OutputFileErr is set by the launcher (via MarkOutputFileUnreliable) when a + // write to OutputFile fails, so the file is known to be incomplete. It is + // empty while the file is trustworthy. When non-empty, neither OutputFile nor + // Result can be treated as the authoritative complete output: the file has a + // gap, and Result is only whatever the worker returned (which may be empty + // while the task is still running, and — depending on the worker — may be a + // partial projection of the file rather than its full content). Consumers + // should report the file's failed state honestly rather than present either + // side as complete. The Manager does not interpret the value beyond emptiness; + // it carries the first failure's message for diagnostics. + OutputFileErr string + // Error contains the error message, set when Status is StatusFailed. + Error string + // RunInBackground indicates whether this task is running (or ran) in the + // background — either launched with RunInBackground, or moved to the background + // after exhausting its foreground budget. It distinguishes background tasks + // from foreground ones when inspecting task state. + RunInBackground bool + // CreatedAt is the time the task was registered. + CreatedAt time.Time + // DoneAt is the time the task reached a terminal state. Nil if still running. + DoneAt *time.Time + // Metadata holds arbitrary extensible fields for future use. + Metadata map[string]any +} + +// TaskEventType describes the lifecycle transition that caused a task event. +type TaskEventType string + +const ( + // TaskEventCreated indicates a task was registered in StatusRunning. + TaskEventCreated TaskEventType = "created" + // TaskEventBackgrounded indicates a foreground task moved to the background. + TaskEventBackgrounded TaskEventType = "backgrounded" + // TaskEventCompleted indicates a task finished successfully. + TaskEventCompleted TaskEventType = "completed" + // TaskEventFailed indicates a task finished with an error. + TaskEventFailed TaskEventType = "failed" + // TaskEventCanceled indicates a task was canceled by Cancel / Close. + TaskEventCanceled TaskEventType = "canceled" +) + +// TaskEvent is a lifecycle event published by Manager.Subscribe. +type TaskEvent struct { + // Type is the transition that caused this event. + Type TaskEventType + // Task is the task snapshot immediately after the transition. + Task *Task +} + +// RunInput is the execution-agnostic input for Run. +// Domain-specific parameters (which agent, which command, the prompt) are +// captured by the WorkFunc closure, not here. +type RunInput struct { + // Description is a short human-readable title for the task, stored in Task.Description. + Description string + // Type is an optional tag for the task (e.g. "bash", "subagent"), stored in + // Task.Type. See Task.Type. + Type string + // ToolUseID is the optional id of the tool call launching this task, stored in + // Task.ToolUseID. See Task.ToolUseID. + ToolUseID string + // RunInBackground controls execution mode: true returns immediately with StatusRunning. + RunInBackground bool + // Metadata is optional caller-supplied data attached to the task's Task.Metadata. + // It is for observers (Get/List, the task_output tool, the host) to correlate or + // label background tasks — e.g. an originating tool-call ID, session, or trace. + // The Manager does not interpret it. It is shallow-copied into the task on creation. + Metadata map[string]any + // OutputFile is the optional path the launcher will write this task's output + // to. When non-empty it is recorded on Task.OutputFile and surfaced in the + // background notice; the Manager itself never writes the file. The launcher + // (a domain adapter) owns writing, so the file may carry interim output while + // the task runs. Empty means the task has no output file. + OutputFile string + // ForegroundTimeoutMs optionally overrides the Manager's foreground budget for + // this run only. When nil, the Manager's configured default applies. When non-nil, + // it bounds how long the run may occupy the foreground before its deadline fires + // (see Config.ShouldAutoBackground for what happens at the deadline). A value <= 0 + // removes the deadline for this run (blocks until completion). Ignored when + // RunInBackground is true. + ForegroundTimeoutMs *int +} + +// defaultForegroundTimeoutMs is the default foreground budget (120 seconds). +const defaultForegroundTimeoutMs = 120_000 + +// IDGenerator returns the complete ID for a new task. +// +// The generator sees the run input before the task is registered and may return a +// business-side identifier. Manager does not add the task-type prefix when IDGen +// is configured; callers that want one should include it in the returned ID. +type IDGenerator func(ctx context.Context, input *RunInput) (string, error) + +// Config configures a Manager. +type Config struct { + // ForegroundTimeoutMs sets the foreground budget: the time a foreground run is + // allowed to occupy the foreground before its deadline fires. + // When > 0, a foreground run that hasn't completed within this many + // milliseconds reaches its deadline (see ShouldAutoBackground for what happens then). + // When 0, there is no deadline (foreground runs block until completion). + // + // Default: 120000ms (120 seconds). + ForegroundTimeoutMs *int + + // ShouldAutoBackground decides, at a foreground run's deadline, whether it may be + // moved to the background (kept running) instead of being canceled. Applications + // can use it to permit long-lived workloads such as servers and watchers while + // timing out commands whose results are no longer useful. The hook receives the + // task, so a host can branch on Task.Type and recover domain parameters from + // Task.Metadata (e.g. the shell command via filesystem.CommandFromTask). + // + // Deciding whether a workload is genuinely long-lived is inherently host- and + // command-specific, so this package ships no built-in policy: the framework + // cannot reliably infer "never exits" from a command string, and a wrong guess + // either kills a useful run or keeps a doomed one. Hosts encode their own rules. + // + // It is consulted ONLY for the auto path — a foreground run that hits its + // deadline. An explicit RunInBackground run always backgrounds immediately, + // regardless of this hook. + // + // When nil (the default), it is treated as always returning false: a run that + // hits its deadline is canceled and reported as timed out, never auto-backgrounded. + ShouldAutoBackground func(ctx context.Context, task *Task) bool + + // IDGen, when set, decides the full ID of every task created by this Manager. + // If nil, Manager uses its default task-type-prefixed base62 ID. + // + // IDGen may be called concurrently by concurrent Run / RunStream calls. It + // must return a non-empty ID. The returned ID must be unique among this + // Manager's registered tasks; a duplicate fails task creation. + IDGen IDGenerator + + // BackgroundNotice customizes the chunk emitted on a RunStream caller's stream + // when a task starts in the background or is auto-moved there. The Manager owns + // only lifecycle facts (id, type, output file); how a host tells the model to + // retrieve the result is host-specific — one host exposes a task_output tool, + // another points at the output file — so that wording does not belong in this + // type-erased layer. + // + // When nil, defaultBackgroundNotice is used: it announces the background launch + // and, when an output file is reserved, directs the reader to Read that path for + // interim output. + // + // The ctx passed to the hook is the run's context (detached from the caller's + // cancellation, carrying its values); use it only for value lookup, not to gate + // the notice on cancellation. + BackgroundNotice func(ctx context.Context, info NoticeInfo) string +} + +// NoticeInfo carries the lifecycle facts a BackgroundNotice hook may use to build +// the chunk shown when a run goes to the background. +type NoticeInfo struct { + // Task is a snapshot of the task at the moment the notice is emitted, carrying + // ID, Type, and OutputFile. Nil only if the task vanished mid-emit (not expected). + Task *Task + // AutoBackgrounded is false when the run was launched directly in the background + // (RunInBackground), and true when a foreground run was auto-moved to the + // background at its deadline because the ShouldAutoBackground hook permitted it + // (a deadline the hook declines becomes a timeout failure, which never reaches + // this notice). The true case is the same transition reported to subscribers as + // TaskEventBackgrounded. + AutoBackgrounded bool +} + +// Manager is a non-generic, in-memory registry that owns the lifecycle of +// managed executions: creation, foreground/background/auto-background +// switching, cancellation and terminal-state tracking. +// +// It is intentionally execution-agnostic: it does not know whether a task is an +// agent or a shell command. Callers launch work via the free function Run, +// passing a WorkFunc that performs the actual execution. A single Manager can +// therefore be shared across multiple domains under one task-ID space. +type Manager struct { + mu sync.Mutex + cond *sync.Cond + tasks map[string]*taskRecord + seq int64 + lastMs int64 + closed bool + foregroundTimeoutMs int + shouldAutoBackground func(ctx context.Context, task *Task) bool + idGen IDGenerator + backgroundNoticeFn func(ctx context.Context, info NoticeInfo) string + + subscribeOnce sync.Once + eventCh chan *TaskEvent + eventBuf *internal.UnboundedChan[*TaskEvent] +} + +type taskRecord struct { + task Task + cancel context.CancelFunc // cancels the run's context + // doneCh is closed exactly once, by finalize, when the task reaches a terminal + // state. Wait selects on it so waiting for one task neither holds m.mu nor is + // woken by unrelated tasks finishing. + doneCh chan struct{} +} + +// New creates a new Manager. +// By default, the foreground budget is 120 seconds; set Config.ForegroundTimeoutMs +// to 0 to remove the deadline (foreground runs block until completion). What +// happens when the budget is reached is governed by Config.ShouldAutoBackground +// (default: cancel the run and report it timed out). +func New(_ context.Context, conf *Config) *Manager { + m := &Manager{ + tasks: make(map[string]*taskRecord), + foregroundTimeoutMs: defaultForegroundTimeoutMs, + } + m.cond = sync.NewCond(&m.mu) + if conf != nil && conf.ForegroundTimeoutMs != nil { + m.foregroundTimeoutMs = *conf.ForegroundTimeoutMs + } + if conf != nil { + m.shouldAutoBackground = conf.ShouldAutoBackground + m.idGen = conf.IDGen + m.backgroundNoticeFn = conf.BackgroundNotice + } + return m +} + +// Subscribe returns a channel that receives TaskEvent values whenever the Manager +// changes a task's lifecycle state. +// +// The stream is forward-only: events generated before the first Subscribe call +// are not replayed (use Get/List to inspect current state). Multiple calls return +// the same shared stream, and Close closes it after buffered events are drained. +// The returned Task values are snapshots; mutating them does not mutate the +// Manager's registry. +func (m *Manager) Subscribe() <-chan *TaskEvent { + m.subscribeOnce.Do(func() { + buf := internal.NewUnboundedChan[*TaskEvent]() + ch := make(chan *TaskEvent) + + m.mu.Lock() + m.eventBuf = buf + m.eventCh = ch + closed := m.closed + m.mu.Unlock() + + go m.relayEvents(buf, ch) + if closed { + buf.Close() + } + }) + return m.eventCh +} + +// relayEvents pumps events from the unbounded buffer to the public channel, +// so publishing under the Manager lock never blocks on a slow subscriber. +func (m *Manager) relayEvents(buf *internal.UnboundedChan[*TaskEvent], ch chan<- *TaskEvent) { + defer close(ch) + for { + event, ok := buf.Receive() + if !ok { + return + } + ch <- event + } +} + +// allowAutoBackground reports whether a run that has hit its foreground deadline +// may be moved to the background. With no configured hook, the answer is false. +func (m *Manager) allowAutoBackground(ctx context.Context, task *Task) bool { + if m.shouldAutoBackground == nil { + return false + } + return m.shouldAutoBackground(ctx, task) +} + +// Get returns the current state of a task by ID. +// Returns (nil, false) if the task does not exist. +func (m *Manager) Get(id string) (*Task, bool) { + m.mu.Lock() + defer m.mu.Unlock() + + rec, ok := m.tasks[id] + if !ok { + return nil, false + } + return cloneTask(&rec.task), true +} + +// Wait blocks until the task with the given id reaches a terminal state, or until +// ctx is canceled, and returns the task's current snapshot together with whether it +// actually reached a terminal state. Callers bound the wait with ctx (e.g. +// context.WithTimeout). +// +// Return values: +// - (nil, false): no task with this id exists. +// - (task, true): the task reached a terminal state (task.Status is terminal). +// - (task, false): ctx was canceled/timed out first; task is the latest +// (still-running) snapshot. +// +// The wait is per-task: it selects on the task's own done channel rather than the +// shared condition, so it neither holds m.mu while waiting nor is woken when other +// tasks finish. +func (m *Manager) Wait(ctx context.Context, id string) (*Task, bool) { + m.mu.Lock() + rec, ok := m.tasks[id] + if !ok { + m.mu.Unlock() + return nil, false + } + doneCh := rec.doneCh + m.mu.Unlock() + + select { + case <-doneCh: + case <-ctx.Done(): + return m.taskSnapshot(id), false + } + return m.taskSnapshot(id), true +} + +// List returns a snapshot of all tasks (both running and completed). +func (m *Manager) List() []*Task { + m.mu.Lock() + defer m.mu.Unlock() + + tasks := make([]*Task, 0, len(m.tasks)) + for _, rec := range m.tasks { + tasks = append(tasks, cloneTask(&rec.task)) + } + return tasks +} + +// Cancel stops a running task. The run's context is canceled and the task +// transitions to StatusCanceled. +// Returns an error if the task does not exist or is not running. +func (m *Manager) Cancel(id string) error { + m.mu.Lock() + defer m.mu.Unlock() + + rec, ok := m.tasks[id] + if !ok { + return fmt.Errorf("no background task has id %q, so there is nothing to stop. "+ + "If you are unsure of the id, there is nothing left to cancel", id) + } + if taskDone(rec.doneCh) { + return fmt.Errorf("background task %q has already finished (status: %s) and cannot be stopped. "+ + "Use the task_output tool with this id to read its result instead", id, rec.task.Status) + } + + m.cancelTask(rec) + return nil +} + +// waitIdle blocks until no registered task is still running, or until the +// provided context is canceled. It backs graceful Close; single-task waits use +// Wait. +func (m *Manager) waitIdle(ctx context.Context) error { + done := make(chan struct{}) + defer close(done) + go func() { + select { + case <-ctx.Done(): + m.cond.Broadcast() + case <-done: + } + }() + + m.mu.Lock() + defer m.mu.Unlock() + + for m.hasRunningLocked() { + if ctx.Err() != nil { + return ctx.Err() + } + m.cond.Wait() + } + return nil +} + +// Close performs graceful shutdown. +// It waits for all running tasks to complete (up to the ctx deadline), +// then cancels any remaining running tasks. +// After Close returns, Run will return an error. +func (m *Manager) Close(ctx context.Context) error { + _ = m.waitIdle(ctx) + + m.mu.Lock() + defer m.mu.Unlock() + + m.closed = true + + for _, rec := range m.tasks { + if !taskDone(rec.doneCh) { + m.cancelTask(rec) + } + } + if m.eventBuf != nil { + m.eventBuf.Close() + } + + return nil +} + +// createTask registers a new task in StatusRunning state. +// The cancel function is not set here — call storeCancelFunc after creation. +func (m *Manager) createTask(ctx context.Context, input *RunInput) (string, error) { + if input == nil { + return "", fmt.Errorf("backgroundtask: RunInput is required") + } + + if m.idGen != nil { + id, err := m.idGen(ctx, input) + if err != nil { + return "", fmt.Errorf("backgroundtask: task id generator: %w", err) + } + return m.registerTask(input, id) + } + + m.mu.Lock() + defer m.mu.Unlock() + + if m.closed { + return "", m.closedError() + } + + id := taskIDPrefix(input.Type) + "_" + base62(m.nextRawID()) + return m.registerTaskLocked(input, id) +} + +func (m *Manager) registerTask(input *RunInput, id string) (string, error) { + m.mu.Lock() + defer m.mu.Unlock() + + if m.closed { + return "", m.closedError() + } + return m.registerTaskLocked(input, id) +} + +func (m *Manager) registerTaskLocked(input *RunInput, id string) (string, error) { + if id == "" { + return "", fmt.Errorf("backgroundtask: task id generator returned empty id") + } + if _, ok := m.tasks[id]; ok { + return "", fmt.Errorf("backgroundtask: task id %q already exists", id) + } + + m.tasks[id] = &taskRecord{ + task: Task{ + ID: id, + Type: input.Type, + ToolUseID: input.ToolUseID, + Description: input.Description, + Status: StatusRunning, + RunInBackground: input.RunInBackground, + CreatedAt: time.Now(), + OutputFile: input.OutputFile, + Metadata: cloneMetadata(input.Metadata), + }, + doneCh: make(chan struct{}), + } + m.sendEventLocked(m.tasks[id], TaskEventCreated) + + return id, nil +} + +func (m *Manager) closedError() error { + return fmt.Errorf("the background task manager has shut down and is no longer accepting new tasks. " + + "Do not retry this; finish using any results you already have") +} + +// cloneMetadata shallow-copies caller-supplied metadata so later mutations to the +// caller's map do not affect the recorded task. Returns nil for empty input. +func cloneMetadata(md map[string]any) map[string]any { + if len(md) == 0 { + return nil + } + clone := make(map[string]any, len(md)) + for k, v := range md { + clone[k] = v + } + return clone +} + +// storeCancelFunc saves the context cancel function for a running task, +// so that Cancel can stop it. +func (m *Manager) storeCancelFunc(id string, cancel context.CancelFunc) { + m.mu.Lock() + defer m.mu.Unlock() + + if rec, ok := m.tasks[id]; ok { + rec.cancel = cancel + } +} + +// MarkOutputFileUnreliable records that a write to the task's output file failed, +// so the file is known to be incomplete. The launcher that owns writing calls it +// (with the task id it receives via WorkFunc's TaskInfo) when an append to the +// file errors. +// +// It sets Task.OutputFileErr on the task with the given id, so consumers stop +// trusting the partial file. The first failure wins: a later call does not +// overwrite an existing message, since once the file has a gap it is unreliable +// regardless of what later writes do. An empty id, an unknown id, or an +// already-marked task is a no-op, so callers may invoke it unconditionally on +// write error. +func (m *Manager) MarkOutputFileUnreliable(taskID, errMsg string) { + if taskID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + + if rec, ok := m.tasks[taskID]; ok && rec.task.OutputFileErr == "" { + rec.task.OutputFileErr = errMsg + } +} + +// completeTask transitions a task to StatusCompleted with the given result. +// No-op if the task is already in a terminal state (idempotent). +func (m *Manager) completeTask(id string, result string) { + m.mu.Lock() + defer m.mu.Unlock() + + m.finalize(id, func(rec *taskRecord) { + rec.task.Status = StatusCompleted + rec.task.Result = result + }) +} + +// failTask transitions a task to StatusFailed with the given error. +// No-op if the task is already in a terminal state (idempotent). +func (m *Manager) failTask(id string, err error) { + m.mu.Lock() + defer m.mu.Unlock() + + m.finalize(id, func(rec *taskRecord) { + rec.task.Status = StatusFailed + if err != nil { + rec.task.Error = err.Error() + } + }) +} + +// timeoutTask transitions a task to StatusFailed with a timed-out error and +// cancels its context (the underlying work is stopped). It is invoked when a +// foreground run hits its deadline and the ShouldAutoBackground hook declined to +// background it. No-op if the task is already terminal (idempotent), so the +// timed-out reason wins the race against the work goroutine's own ctx-canceled error. +func (m *Manager) timeoutTask(id string, budgetMs int) { + m.mu.Lock() + defer m.mu.Unlock() + + rec, ok := m.tasks[id] + if !ok || taskDone(rec.doneCh) { + return + } + if rec.cancel != nil { + rec.cancel() + } + m.finalize(id, func(rec *taskRecord) { + rec.task.Status = StatusFailed + rec.task.Error = fmt.Sprintf("timed out after %dms", budgetMs) + }) +} + +// canceledError is the message recorded on Task.Error when a task is stopped by +// Cancel or Close, so the canceled outcome carries a reason rather than an empty +// terminal state. +const canceledError = "task was canceled" + +// cancelTask transitions a task to StatusCanceled and cancels its context. +// Must be called with m.mu held. +func (m *Manager) cancelTask(rec *taskRecord) { + if rec.cancel != nil { + rec.cancel() + } + + m.finalize(rec.task.ID, func(rec *taskRecord) { + rec.task.Status = StatusCanceled + rec.task.Error = canceledError + }) +} + +// cancelIfRunning cancels a task's work and marks it StatusCanceled if it is +// still running. Used when a foreground caller abandons its wait (its context is +// canceled before the task completes or auto-backgrounds). Idempotent: a no-op if +// the task already reached a terminal state. +func (m *Manager) cancelIfRunning(id string) { + m.mu.Lock() + defer m.mu.Unlock() + + rec, ok := m.tasks[id] + if !ok || taskDone(rec.doneCh) { + return + } + m.cancelTask(rec) +} + +// detach marks a still-running task as a background task and publishes an event. +// It is called when Run hands the task back to the caller as StatusRunning (an +// explicit background launch, or an auto-background at the foreground deadline). +// +// Returns false if the task already reached a terminal state — i.e. the work +// finished concurrently with the deadline — in which case the caller should report +// the actual outcome instead of StatusRunning. +func (m *Manager) detach(id string) bool { + m.mu.Lock() + defer m.mu.Unlock() + + rec, ok := m.tasks[id] + if !ok || taskDone(rec.doneCh) { + return false + } + rec.task.RunInBackground = true + m.sendEventLocked(rec, TaskEventBackgrounded) + return true +} + +// finalize applies a terminal state transition to a task. +// It sets done=true, records DoneAt, publishes an event, and broadcasts the +// condition. +// Returns false if the task was not found or already in a terminal state (idempotent). +// Must be called with m.mu held. +func (m *Manager) finalize(id string, apply func(rec *taskRecord)) bool { + rec, ok := m.tasks[id] + if !ok || taskDone(rec.doneCh) { + return false + } + + now := time.Now() + rec.task.DoneAt = &now + apply(rec) + + // Signal per-task waiters (Wait) and the all-done waiters (Close). + close(rec.doneCh) + m.sendEventLocked(rec, eventTypeForStatus(rec.task.Status)) + m.cond.Broadcast() + return true +} + +// sendEventLocked publishes a task event to subscribers, if Subscribe has +// been called. Must be called with m.mu held. +func (m *Manager) sendEventLocked(rec *taskRecord, typ TaskEventType) { + if typ != "" && m.eventBuf != nil { + m.eventBuf.TrySend(&TaskEvent{Type: typ, Task: cloneTask(&rec.task)}) + } +} + +func eventTypeForStatus(status Status) TaskEventType { + switch status { + case StatusCompleted: + return TaskEventCompleted + case StatusFailed: + return TaskEventFailed + case StatusCanceled: + return TaskEventCanceled + default: + return "" + } +} + +// cloneTask returns a copy of t safe to hand to callers. The Metadata map is +// shallow-copied so callers cannot mutate the registry's map entries, though +// mutable values stored inside Metadata remain shared. +func cloneTask(t *Task) *Task { + clone := *t + clone.Metadata = cloneMetadata(t.Metadata) + return &clone +} + +func (m *Manager) hasRunningLocked() bool { + for _, rec := range m.tasks { + if !taskDone(rec.doneCh) { + return true + } + } + return false +} + +func taskDone(doneCh <-chan struct{}) bool { + select { + case <-doneCh: + return true + default: + return false + } +} + +// TaskInfo is a read-only snapshot of the facts the Manager establishes about a +// task at creation, handed to the WorkFunc when it starts. It is not the live +// Task record: it carries only identity fixed at creation, never the mutable +// lifecycle fields (Result/Status/Error/OutputFileErr) the Manager fills later, +// so work never races on them. The work already holds everything the launcher +// passed in (Type, OutputFile, Metadata, ...); TaskInfo supplies what only the +// Manager knows. New fields may be added over time — adding a field is backward +// compatible, so the WorkFunc signature stays stable. +type TaskInfo struct { + // ID is the Manager-generated task id. It is the one fact the work cannot + // otherwise obtain: the id is assigned inside createTask, after the work + // closure is already built. The launcher passes it to MarkOutputFileUnreliable + // to report an output-file write failure against this task. + ID string +} + +// WorkFunc performs a single managed execution. It is supplied by the caller +// (e.g. a subagent or filesystem adapter); the Manager itself never knows what +// the work is. +// +// task carries the Manager-assigned facts about this run (see TaskInfo) — most +// importantly its id, which the work needs to report an output-file write failure +// via Manager.MarkOutputFileUnreliable. +// +// ctx carries the values of the Run call's context but is detached from its +// cancellation, so a backgrounded task outlives the turn that launched it. It is +// canceled when Cancel is invoked for this task, when a foreground deadline or an +// abandoned foreground wait stops it, or when the Manager is closed. Work should +// honor it. +// +// The returned result becomes Task.Result; a non-nil err becomes Task.Error and +// transitions the task to StatusFailed. +type WorkFunc func(ctx context.Context, task TaskInfo) (result string, err error) + +// detachedCtx carries its parent's values but is never canceled by the parent. +// It mirrors context.WithoutCancel (Go 1.21+); this package targets Go 1.18. +// Background work runs under a detachedCtx (wrapped by a fresh cancelable context) +// so it survives cancellation of the per-turn context that launched it, while +// still seeing that context's values. +type detachedCtx struct{ parent context.Context } + +func (detachedCtx) Deadline() (deadline time.Time, ok bool) { return time.Time{}, false } + +func (detachedCtx) Done() <-chan struct{} { return nil } + +func (detachedCtx) Err() error { return nil } + +func (c detachedCtx) Value(key any) any { return c.parent.Value(key) } + +// Run executes work as a managed task on m. +// +// The execution mode depends on input.RunInBackground and the effective foreground +// budget (input.ForegroundTimeoutMs if set, else the Manager's configured default): +// - Foreground (RunInBackground=false, budget<=0): blocks until completion +// - Background (RunInBackground=true): returns immediately with StatusRunning +// - Deadline (budget>0): runs in foreground up to the budget, then — if still +// running — consults the Manager's ShouldAutoBackground hook. If it permits, +// the run is moved to the background (kept running) and Run returns +// StatusRunning. Otherwise the run is canceled and reported as timed out +// (StatusFailed). +// +// All runs are tracked in Manager state and visible via Get/List. +func (m *Manager) Run(ctx context.Context, input *RunInput, work WorkFunc) (*Task, error) { + id, err := m.createTask(ctx, input) + if err != nil { + return nil, err + } + + // The work runs under a context detached from the caller's (per-turn) + // cancellation, so a backgrounded task is not killed when the turn that + // launched it ends or is preempted. The caller ctx's values are preserved + // (framework/session state the work relies on); only its cancellation is + // dropped. The work is stopped by Cancel(id), the foreground deadline, an + // abandoned foreground wait (caller ctx canceled), or Close. + runCtx, cancel := context.WithCancel(detachedCtx{parent: ctx}) + m.storeCancelFunc(id, cancel) + + // run executes the work and finalizes the task. The terminal outcome lives on + // the task record (set by completeTask/failTask), which Run reads back via + // taskSnapshot — so run signals completion rather than returning a value. + run := func() { + defer cancel() + defer func() { + if p := recover(); p != nil { + // A panicking WorkFunc must fail its own task, not crash the process. + m.failTask(id, safe.NewPanicErr(p, debug.Stack())) + } + }() + r, runErr := work(runCtx, TaskInfo{ID: id}) + if runErr != nil { + m.failTask(id, runErr) + } else { + m.completeTask(id, r) + } + } + + // Explicit background: run in goroutine, return immediately. createTask already + // marked the task RunInBackground. + if input.RunInBackground { + go run() + return m.taskSnapshot(id), nil + } + + // Foreground: run in a goroutine and wait. The wait honors caller cancellation + // (the detached work ctx does not, so it is canceled explicitly here) and, when a + // budget is set, the foreground deadline. + done := make(chan struct{}, 1) + go func() { run(); done <- struct{}{} }() + + budgetMs := m.foregroundTimeoutMs + if input.ForegroundTimeoutMs != nil { + budgetMs = *input.ForegroundTimeoutMs + } + + if budgetMs > 0 { + // Foreground with a deadline: wait up to the effective budget (per-run + // override takes precedence over the Manager default). On the deadline, + // either move to the background (if the hook permits) or cancel as timed out. + timer := time.NewTimer(time.Duration(budgetMs) * time.Millisecond) + defer timer.Stop() + select { + case <-done: + // run() has already finalized the task before signaling, so the recorded + // state is authoritative — including a StatusCanceled set by a concurrent + // Cancel, which must win over the work's own ctx-canceled error. + return m.taskSnapshot(id), nil + case <-ctx.Done(): + // Caller abandoned the foreground wait before the deadline (e.g. the + // turn was canceled) and before any auto-background: stop the work. + m.cancelIfRunning(id) + return m.taskSnapshot(id), nil + case <-timer.C: + task, ok := m.Get(id) + if !ok || task.DoneAt != nil { + // Work finished right at the deadline — report its actual outcome. + return m.taskSnapshot(id), nil + } + if m.allowAutoBackground(ctx, task) && m.detach(id) { + return m.taskSnapshot(id), nil + } + // Hook declined (or work finished during the hook): stop if still running. + m.timeoutTask(id, budgetMs) + return m.taskSnapshot(id), nil + } + } + + // Foreground without a deadline: block until completion or caller cancellation. + select { + case <-done: + return m.taskSnapshot(id), nil + case <-ctx.Done(): + m.cancelIfRunning(id) + return m.taskSnapshot(id), nil + } +} + +// StreamWorkFunc performs a single managed streaming execution. It is the +// streaming counterpart of WorkFunc: instead of returning the whole result at +// once, it returns a stream of output chunks. The Manager forwards those chunks +// to the RunStream caller in real time and, in parallel, accumulates them into +// the task's final Result (and OutputFile). Chunk semantics (formatting, exit +// codes) are entirely the caller's concern; the Manager only concatenates. +// +// task behaves exactly as for WorkFunc (see TaskInfo): it carries the task id the +// work uses to report an output-file write failure. +// +// ctx behaves exactly as for WorkFunc (see WorkFunc): detached from the caller's +// cancellation, stopped by Cancel/deadline/Close. Work should honor it and close +// the returned reader when ctx is done. +type StreamWorkFunc func(ctx context.Context, task TaskInfo) (*schema.StreamReader[string], error) + +// RunStream executes streaming work as a managed task, returning a stream of +// output chunks to consume in real time. +// +// It mirrors Run's lifecycle (tracking, foreground budget, auto-background) but +// preserves streaming for the foreground phase: +// - Foreground completion: every chunk is forwarded live, then the stream closes. +// - Auto-background at the deadline: chunks forwarded so far are kept; the +// Manager appends a single notice chunk (task id + output file) and closes the +// caller's stream, while the work keeps running in the background — its +// remaining output is drained into the task's Result/OutputFile. +// - Explicit background (input.RunInBackground): the work runs detached from the +// start, so no execution chunks reach the caller; the stream carries only the +// background notice and then closes. +// +// The returned reader is always non-nil on a nil error. The Manager is the sole +// writer of that stream, so there is never a write race with the work. +func (m *Manager) RunStream(ctx context.Context, input *RunInput, work StreamWorkFunc) (*schema.StreamReader[string], error) { + id, err := m.createTask(ctx, input) + if err != nil { + return nil, err + } + + runCtx, cancel := context.WithCancel(detachedCtx{parent: ctx}) + m.storeCancelFunc(id, cancel) + + sr, sw := schema.Pipe[string](streamBufferCap) + + budgetMs := m.foregroundTimeoutMs + if input.ForegroundTimeoutMs != nil { + budgetMs = *input.ForegroundTimeoutMs + } + // An explicit background launch has no foreground phase to stream, so its + // budget is irrelevant: forward nothing, just emit the notice. + if input.RunInBackground { + budgetMs = 0 + } + + go m.forwardStream(&streamRun{ + callerCtx: ctx, + runCtx: runCtx, + cancel: cancel, + id: id, + input: input, + work: work, + sw: sw, + budgetMs: budgetMs, + }) + return sr, nil +} + +// streamRun bundles the per-run state for forwardStream (kept as one value to stay +// within the argument limit and to make the goroutine launch self-documenting). +type streamRun struct { + callerCtx context.Context + runCtx context.Context + cancel context.CancelFunc + id string + input *RunInput + work StreamWorkFunc + sw *schema.StreamWriter[string] + budgetMs int +} + +// forwardStream owns the caller-facing stream writer sw: it is the only goroutine +// that writes to it, so injecting the background notice never races the work. +func (m *Manager) forwardStream(r *streamRun) { + defer r.cancel() + // A panic constructing the work stream (r.work below) lands here, before sw is + // closed: fail the task and surface it on the caller stream. Panics while reading + // chunks are recovered closer to their source — pumpStream for the foreground + // loop, drainReader for the background drain — so this never double-closes sw. + defer func() { + if p := recover(); p != nil { + err := safe.NewPanicErr(p, debug.Stack()) + m.failTask(r.id, err) + r.sw.Send("", err) + r.sw.Close() + } + }() + + ws, err := r.work(r.runCtx, TaskInfo{ID: r.id}) + if err != nil { + m.failTask(r.id, err) + r.sw.Send("", err) + r.sw.Close() + return + } + defer ws.Close() + + var buf strings.Builder + + // Explicit background: no foreground phase to stream and no deadline to race + // (RunStream forces budgetMs=0), so skip the pumpStream goroutine entirely. + // Emit the notice, close the caller stream, and drain the reader directly into + // the result. + if r.input.RunInBackground { + r.sw.Send(m.backgroundStartNotice(r.runCtx, r.id), nil) + r.sw.Close() + m.drainReader(r.id, ws, &buf) + return + } + + chunks := pumpStream(r.runCtx, ws) + + var timerC <-chan time.Time + if r.budgetMs > 0 { + timer := time.NewTimer(time.Duration(r.budgetMs) * time.Millisecond) + defer timer.Stop() + timerC = timer.C + } + + for { + select { + case c := <-chunks: + if c.err == io.EOF { + m.completeTask(r.id, buf.String()) + r.sw.Close() + return + } + if c.err != nil { + m.failTask(r.id, c.err) + r.sw.Send("", c.err) + r.sw.Close() + return + } + buf.WriteString(c.text) + if r.sw.Send(c.text, nil) { + // Caller closed the stream early (abandoned the read): stop the work. + m.cancelIfRunning(r.id) + return + } + case <-r.callerCtx.Done(): + // Caller abandoned the foreground wait before the deadline: stop work. + m.cancelIfRunning(r.id) + r.sw.Close() + return + case <-timerC: + task, ok := m.Get(r.id) + if !ok || task.DoneAt != nil { + continue // finished right at the deadline; let the chunks case end it + } + if m.allowAutoBackground(r.callerCtx, task) && m.detach(r.id) { + // Moved to the background: cap the caller's stream with a notice and + // keep draining the rest into the task result. + r.sw.Send(m.backgroundMoveNotice(r.runCtx, r.id), nil) + r.sw.Close() + m.drainStream(r.runCtx, r.id, chunks, &buf) + return + } + m.timeoutTask(r.id, r.budgetMs) + r.sw.Close() + return + } + } +} + +// streamChunk is one item pumped off a StreamWorkFunc's reader: either a piece of +// output text or a terminal error (io.EOF on normal completion). +type streamChunk struct { + text string + err error +} + +// pumpStream turns a stream reader's blocking Recv loop into a channel, so the +// forward loop can wait on output alongside the deadline and caller cancellation. +// It stops when ctx is done (the run was canceled, timed out, or closed). +func pumpStream(ctx context.Context, ws *schema.StreamReader[string]) <-chan streamChunk { + chunks := make(chan streamChunk) + go func() { + // ws.Recv runs the work's stream (including any convert step) in this + // goroutine, so a panic there must not crash the process. Turn it into a + // terminal chunk error; the forward loop fails the task on it like any other. + defer func() { + if p := recover(); p != nil { + select { + case chunks <- streamChunk{err: safe.NewPanicErr(p, debug.Stack())}: + case <-ctx.Done(): + } + } + }() + for { + text, err := ws.Recv() + c := streamChunk{text: text, err: err} + select { + case chunks <- c: + case <-ctx.Done(): + return + } + if err != nil { + return + } + } + }() + return chunks +} + +// drainReader consumes a backgrounded run's reader directly into buf and finalizes +// the task on completion. Used by the explicit-background path, which has no +// foreground select loop and therefore no pumpStream channel — reading ws inline +// saves a goroutine. Called after the caller's stream has been closed, so it never +// writes to sw. +// +// Unlike drainStream it has no runCtx.Done() case: it exits only when ws.Recv +// returns EOF or an error. On Cancel/Close the run's ctx is canceled, which the +// work must honor by ending its stream — that is what unblocks Recv here. A no-op +// completeTask/failTask then loses the race against the cancelTask that already +// finalized the task (both are idempotent). This mirrors the original drainStream, +// whose pump goroutine likewise stayed blocked on Recv until the work honored ctx. +func (m *Manager) drainReader(id string, ws *schema.StreamReader[string], buf *strings.Builder) { + // ws.Recv runs the work's stream in this goroutine; a panic must fail the task, + // not crash the process. The caller stream is already closed before draining, so + // recovery only finalizes — it never touches sw. + defer func() { + if p := recover(); p != nil { + m.failTask(id, safe.NewPanicErr(p, debug.Stack())) + } + }() + for { + text, err := ws.Recv() + if err == io.EOF { + m.completeTask(id, buf.String()) + return + } + if err != nil { + m.failTask(id, err) + return + } + buf.WriteString(text) + } +} + +// drainStream consumes a backgrounded run's remaining chunks into buf and +// finalizes the task on completion. Called after the caller's stream has been +// closed, so it never writes to sw. +func (m *Manager) drainStream(runCtx context.Context, id string, chunks <-chan streamChunk, buf *strings.Builder) { + for { + select { + case c := <-chunks: + if c.err == io.EOF { + m.completeTask(id, buf.String()) + return + } + if c.err != nil { + m.failTask(id, c.err) + return + } + buf.WriteString(c.text) + case <-runCtx.Done(): + // The run was canceled/closed while backgrounded; finalize already + // happened via cancelTask, so just stop draining. + return + } + } +} + +// backgroundStartNotice builds the chunk emitted for an explicit RunInBackground +// launch. +func (m *Manager) backgroundStartNotice(ctx context.Context, id string) string { + return m.notice(ctx, id, false) +} + +// backgroundMoveNotice builds the chunk appended when a foreground run is moved to +// the background by the auto-background policy. +func (m *Manager) backgroundMoveNotice(ctx context.Context, id string) string { + return m.notice(ctx, id, true) +} + +// notice produces the background-launch chunk: the configured BackgroundNotice +// hook when set, otherwise defaultBackgroundNotice. It snapshots the task so the +// hook sees the same lifecycle facts (id, type, output file) the default would. +func (m *Manager) notice(ctx context.Context, id string, autoBackgrounded bool) string { + task, _ := m.Get(id) + info := NoticeInfo{Task: task, AutoBackgrounded: autoBackgrounded} + if m.backgroundNoticeFn != nil { + return m.backgroundNoticeFn(ctx, info) + } + return defaultBackgroundNotice(info) +} + +// noticeTemplate is the default background-notice text. Placeholders are filled by +// defaultBackgroundNotice; {kind} and {output} expand to empty when absent, so the +// same template serves the with- and without-output-file cases. +const noticeTemplate = "\n[task {id}{kind} {state}; you will be notified when it completes.{output}]" + +// noticeOutputTemplate is the {output} fragment, present only when the task has a +// reserved output file. +const noticeOutputTemplate = " Output is being written to: {file}." + + " To check interim output, use Read on that file path." + +// defaultBackgroundNotice is the built-in BackgroundNotice. It announces the +// background launch and, when an output file is reserved, directs the reader to +// Read that path for interim output. It deliberately names no control tool, since +// the retrieval mechanism is host-specific (see Config.BackgroundNotice). +func defaultBackgroundNotice(info NoticeInfo) string { + id, kind, outputFile := "", "", "" + if info.Task != nil { + id = info.Task.ID + if info.Task.Type != "" { + kind = " (" + info.Task.Type + ")" + } + outputFile = info.Task.OutputFile + } + + state := "is running in the background" + if info.AutoBackgrounded { + state = "moved to the background" + } + + output := "" + if outputFile != "" { + output = strings.NewReplacer("{file}", outputFile).Replace(noticeOutputTemplate) + } + + return strings.NewReplacer( + "{id}", id, + "{kind}", kind, + "{state}", state, + "{output}", output, + ).Replace(noticeTemplate) +} + +// streamBufferCap is the buffer size of the caller-facing stream pipe. +const streamBufferCap = 16 + +// taskSnapshot returns the current state of a task as a cloned *Task. The task +// record is the single source of truth, so a concurrent cancel or timeout is +// reflected faithfully. It falls back to a minimal failed snapshot if the task is +// somehow absent (should not happen for a just-created task). +func (m *Manager) taskSnapshot(id string) *Task { + if task, ok := m.Get(id); ok { + return task + } + return &Task{ID: id, Status: StatusFailed} +} diff --git a/adk/backgroundtask/manager_test.go b/adk/backgroundtask/manager_test.go new file mode 100644 index 000000000..c39f7b944 --- /dev/null +++ b/adk/backgroundtask/manager_test.go @@ -0,0 +1,791 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package backgroundtask + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// closeWithTimeout closes the Manager with a short timeout to avoid blocking on uncompleted tasks. +func closeWithTimeout(m *Manager) { + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + _ = m.Close(ctx) +} + +func intPtr(v int) *int { return &v } + +// anyRunning reports whether the manager still has a task in StatusRunning, +// derived from the public List() snapshot. +func anyRunning(m *Manager) bool { + for _, t := range m.List() { + if t.Status == StatusRunning { + return true + } + } + return false +} + +// workReturning builds a WorkFunc that returns the given result/error immediately. +func workReturning(result string, err error) WorkFunc { + return func(ctx context.Context, _ TaskInfo) (string, error) { + return result, err + } +} + +// workSleeping builds a WorkFunc that sleeps then returns result. +func workSleeping(d time.Duration, result string) WorkFunc { + return func(ctx context.Context, _ TaskInfo) (string, error) { + time.Sleep(d) + return result, nil + } +} + +// workBlocking builds a WorkFunc that blocks until its context is canceled. +func workBlocking() WorkFunc { + return func(ctx context.Context, _ TaskInfo) (string, error) { + <-ctx.Done() + return "", ctx.Err() + } +} + +func run(m *Manager, description string, background bool, work WorkFunc) (*Task, error) { + return m.Run(context.Background(), &RunInput{ + Description: description, + RunInBackground: background, + }, work) +} + +func waitTask(t *testing.T, m *Manager, id string) *Task { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + task, done := m.Wait(ctx, id) + require.NotNil(t, task) + require.True(t, done, "task %s did not finish before the wait deadline", id) + return task +} + +func waitTaskEvent(t *testing.T, ch <-chan *TaskEvent, match func(*TaskEvent) bool) *TaskEvent { + t.Helper() + timeout := time.After(time.Second) + for { + select { + case event, ok := <-ch: + require.True(t, ok, "subscription closed before the expected update") + if match(event) { + return event + } + case <-timeout: + t.Fatal("timed out waiting for the expected task update") + } + } +} + +// --- Run (foreground) Tests --- + +func TestManager_RunForeground(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "test task", false, workReturning("hello", nil)) + require.NoError(t, err) + assert.Equal(t, StatusCompleted, result.Status) + assert.Equal(t, "hello", result.Result) + assert.NotEmpty(t, result.ID) +} + +func TestManager_RunForegroundError(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "failing task", false, workReturning("", fmt.Errorf("something failed"))) + require.NoError(t, err) // Run itself doesn't error + assert.Equal(t, StatusFailed, result.Status) + assert.Equal(t, "something failed", result.Error) +} + +// --- Run (background) Tests --- + +func TestManager_RunBackground(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "bg task", true, workSleeping(50*time.Millisecond, "bg result")) + require.NoError(t, err) + assert.Equal(t, StatusRunning, result.Status) + assert.NotEmpty(t, result.ID) + assert.True(t, anyRunning(m)) + + task := waitTask(t, m, result.ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "bg result", task.Result) +} + +// --- Work context lifetime Tests --- + +type bgCtxKey string + +// A backgrounded task must survive cancellation of the per-call (per-turn) +// context that launched it: it is stopped only by Cancel/Close/deadline. +func TestManager_RunBackground_SurvivesCallerCtxCancel(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + callerCtx, cancelCaller := context.WithCancel(context.Background()) + + started := make(chan struct{}) + release := make(chan struct{}) + result, err := m.Run(callerCtx, &RunInput{Description: "bg", RunInBackground: true}, + func(ctx context.Context, _ TaskInfo) (string, error) { + close(started) + select { + case <-release: + return "done", nil + case <-ctx.Done(): + return "", ctx.Err() + } + }) + require.NoError(t, err) + require.Equal(t, StatusRunning, result.Status) + <-started + + // Cancel the caller (per-turn) context; the background task must keep running. + cancelCaller() + time.Sleep(50 * time.Millisecond) + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, StatusRunning, task.Status, "background task should survive caller ctx cancellation") + + // It finishes only when the work itself completes. + close(release) + task = waitTask(t, m, result.ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "done", task.Result) +} + +// A foreground task with no deadline must still be stopped when the caller +// abandons its wait (per-call context canceled). +func TestManager_RunForeground_CallerCtxCancelStops(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(0)}) + defer closeWithTimeout(m) + + callerCtx, cancelCaller := context.WithCancel(context.Background()) + go func() { + time.Sleep(30 * time.Millisecond) + cancelCaller() + }() + + result, err := m.Run(callerCtx, &RunInput{Description: "fg blocking"}, workBlocking()) + require.NoError(t, err) + assert.Equal(t, StatusCanceled, result.Status) +} + +// The work context preserves the caller context's values (framework/session +// state) even though it is detached from the caller's cancellation. +func TestManager_RunBackground_PreservesCallerCtxValues(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + const key bgCtxKey = "trace" + callerCtx := context.WithValue(context.Background(), key, "abc") + + got := make(chan interface{}, 1) + result, err := m.Run(callerCtx, &RunInput{Description: "bg", RunInBackground: true}, + func(ctx context.Context, _ TaskInfo) (string, error) { + got <- ctx.Value(key) + return "ok", nil + }) + require.NoError(t, err) + require.Equal(t, StatusRunning, result.Status) + waitTask(t, m, result.ID) + + select { + case v := <-got: + assert.Equal(t, "abc", v, "background work should see caller ctx values") + case <-time.After(time.Second): + t.Fatal("work did not run") + } +} + +// --- Subscribe Tests --- + +func TestManager_Subscribe_ForegroundLifecycle(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + ch := m.Subscribe() + result, err := run(m, "fg task", false, workReturning("done", nil)) + require.NoError(t, err) + + created := waitTaskEvent(t, ch, func(event *TaskEvent) bool { + return event.Type == TaskEventCreated && event.Task.ID == result.ID + }) + assert.False(t, created.Task.RunInBackground) + assert.Equal(t, StatusRunning, created.Task.Status) + + completed := waitTaskEvent(t, ch, func(event *TaskEvent) bool { + return event.Type == TaskEventCompleted && event.Task.ID == result.ID + }) + assert.Equal(t, StatusCompleted, completed.Task.Status) + assert.Equal(t, "done", completed.Task.Result) +} + +func TestManager_Subscribe_BackgroundLifecycle(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + ch := m.Subscribe() + result, err := run(m, "bg task", true, workSleeping(20*time.Millisecond, "bg result")) + require.NoError(t, err) + + created := waitTaskEvent(t, ch, func(event *TaskEvent) bool { + return event.Type == TaskEventCreated && event.Task.ID == result.ID + }) + assert.True(t, created.Task.RunInBackground) + assert.Equal(t, StatusRunning, created.Task.Status) + + done := waitTaskEvent(t, ch, func(event *TaskEvent) bool { + return event.Type == TaskEventCompleted && event.Task.ID == result.ID + }) + assert.Equal(t, StatusCompleted, done.Task.Status) + assert.Equal(t, "bg result", done.Task.Result) + assert.NotNil(t, done.Task.DoneAt) +} + +func TestManager_Subscribe_AutoBackgroundChange(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(20), ShouldAutoBackground: allowBackground}) + defer closeWithTimeout(m) + + ch := m.Subscribe() + result, err := run(m, "slow", false, workSleeping(80*time.Millisecond, "late")) + require.NoError(t, err) + assert.Equal(t, StatusRunning, result.Status) + + bg := waitTaskEvent(t, ch, func(event *TaskEvent) bool { + return event.Type == TaskEventBackgrounded && event.Task.ID == result.ID + }) + assert.Equal(t, StatusRunning, bg.Task.Status) + assert.True(t, bg.Task.RunInBackground) + assert.Equal(t, "slow", bg.Task.Description) + + done := waitTaskEvent(t, ch, func(event *TaskEvent) bool { + return event.Type == TaskEventCompleted && event.Task.ID == result.ID + }) + assert.Equal(t, "late", done.Task.Result) +} + +func TestManager_Subscribe_CancelChange(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + ch := m.Subscribe() + result, err := run(m, "bg", true, workBlocking()) + require.NoError(t, err) + require.NoError(t, m.Cancel(result.ID)) + + done := waitTaskEvent(t, ch, func(event *TaskEvent) bool { + return event.Type == TaskEventCanceled && event.Task.ID == result.ID + }) + assert.Equal(t, canceledError, done.Task.Error) +} + +func TestManager_Subscribe_ClosesOnClose(t *testing.T) { + m := New(context.Background(), &Config{}) + ch := m.Subscribe() + + require.NoError(t, m.Close(context.Background())) + _, ok := <-ch + assert.False(t, ok) +} + +// --- Type / ToolUseID --- + +func TestManager_TypeAndToolUseIDStored(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := m.Run(context.Background(), &RunInput{ + Description: "task", + Type: "bash", + ToolUseID: "call_42", + }, workReturning("done", nil)) + require.NoError(t, err) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, "bash", task.Type) + assert.Equal(t, "call_42", task.ToolUseID) +} + +// --- Output file --- + +// The Manager records RunInput.OutputFile on the task and surfaces it, but never +// writes the file itself (the launcher owns writing). +func TestManager_OutputFile_RecordedNotWritten(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := m.Run(context.Background(), &RunInput{ + Description: "task", + OutputFile: "/tasks/custom.output", + }, workReturning("the output", nil)) + require.NoError(t, err) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, "/tasks/custom.output", task.OutputFile) + // Result is still tracked in memory; the Manager does not touch the file. + assert.Equal(t, "the output", task.Result) +} + +func TestManager_NoOutputFile(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "task", false, workReturning("the output", nil)) + require.NoError(t, err) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Empty(t, task.OutputFile) + assert.Equal(t, "the output", task.Result) +} + +// --- Auto-background Tests --- + +// allowBackground is a ShouldAutoBackground hook that permits backgrounding any run. +func allowBackground(context.Context, *Task) bool { return true } + +func TestManager_AutoBackground_Slow(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(50), ShouldAutoBackground: allowBackground}) + defer closeWithTimeout(m) + + result, err := run(m, "slow task", false, workSleeping(200*time.Millisecond, "slow result")) + require.NoError(t, err) + assert.Equal(t, StatusRunning, result.Status) + assert.True(t, anyRunning(m)) + + task := waitTask(t, m, result.ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "slow result", task.Result) +} + +// A per-run ForegroundTimeoutMs overrides the Manager default: here the Manager has +// auto-background disabled (0), but the run sets a short per-call deadline, so a +// slow command is moved to the background (the hook permits it) rather than blocking. +func TestManager_PerRunAutoBackgroundOverride(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(0), ShouldAutoBackground: allowBackground}) + defer closeWithTimeout(m) + + override := 50 + result, err := m.Run(context.Background(), &RunInput{ + Description: "slow", + ForegroundTimeoutMs: &override, + }, workSleeping(300*time.Millisecond, "slow result")) + require.NoError(t, err) + assert.Equal(t, StatusRunning, result.Status) // moved to background at 50ms + assert.True(t, anyRunning(m)) + + task := waitTask(t, m, result.ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "slow result", task.Result) +} + +// With no ShouldAutoBackground hook (the default), a run that hits its deadline is +// canceled and reported as timed out — not backgrounded. +func TestManager_DeadlineKillsWhenNotBackgroundable(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(50)}) // no hook + defer closeWithTimeout(m) + + result, err := run(m, "slow task", false, workBlocking()) + require.NoError(t, err) + assert.Equal(t, StatusFailed, result.Status) + assert.Contains(t, result.Error, "timed out") + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, StatusFailed, task.Status) + assert.False(t, anyRunning(m)) // work was canceled +} + +// The hook receives the task so the business can decide per-run; here it backgrounds +// only tasks whose description marks them as a server. +func TestManager_ShouldAutoBackgroundPerTask(t *testing.T) { + m := New(context.Background(), &Config{ + ForegroundTimeoutMs: intPtr(40), + ShouldAutoBackground: func(_ context.Context, task *Task) bool { + return task.Description == "server" + }, + }) + defer closeWithTimeout(m) + + bg, err := run(m, "server", false, workSleeping(150*time.Millisecond, "up")) + require.NoError(t, err) + assert.Equal(t, StatusRunning, bg.Status) // backgrounded + + killed, err := run(m, "oneshot", false, workBlocking()) + require.NoError(t, err) + assert.Equal(t, StatusFailed, killed.Status) // timed out + assert.Contains(t, killed.Error, "timed out") + + waitTask(t, m, bg.ID) +} + +// A per-run override of <=0 disables auto-background even when the Manager has a +// default, so the run blocks until completion. +func TestManager_PerRunAutoBackgroundDisable(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(20)}) // would auto-bg fast + defer closeWithTimeout(m) + + off := 0 + result, err := m.Run(context.Background(), &RunInput{ + Description: "blocking-foreground", + ForegroundTimeoutMs: &off, + }, workSleeping(60*time.Millisecond, "done")) + require.NoError(t, err) + assert.Equal(t, StatusCompleted, result.Status) // blocked despite the 20ms default + assert.Equal(t, "done", result.Result) +} + +func TestManager_AutoBackground_Fast(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(5000)}) + defer closeWithTimeout(m) + + result, err := run(m, "fast task", false, workReturning("fast result", nil)) + require.NoError(t, err) + assert.Equal(t, StatusCompleted, result.Status) + assert.Equal(t, "fast result", result.Result) + assert.False(t, anyRunning(m)) +} + +// --- Get/List Tests --- + +func TestManager_GetNotFound(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + task, ok := m.Get("nonexistent") + assert.False(t, ok) + assert.Nil(t, task) +} + +func TestManager_Get(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "test task", false, workReturning("done", nil)) + require.NoError(t, err) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, result.ID, task.ID) + assert.Equal(t, "test task", task.Description) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "done", task.Result) + assert.NotNil(t, task.DoneAt) +} + +func TestManager_Metadata(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + md := map[string]any{"toolCallID": "call_42", "session": "s1"} + result, err := m.Run(context.Background(), &RunInput{ + Description: "task", + Metadata: md, + }, workReturning("done", nil)) + require.NoError(t, err) + + // Metadata flows to the tracked task, visible via Get. + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, "call_42", task.Metadata["toolCallID"]) + assert.Equal(t, "s1", task.Metadata["session"]) + + // Mutating the caller's original map must not affect the recorded task. + md["toolCallID"] = "mutated" + task, _ = m.Get(result.ID) + assert.Equal(t, "call_42", task.Metadata["toolCallID"]) +} + +func TestManager_List(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + r1, _ := run(m, "task1", false, workReturning("r1", nil)) + r2, _ := run(m, "task2", false, workReturning("r2", nil)) + + tasks := m.List() + assert.Len(t, tasks, 2) + + byID := make(map[string]*Task) + for _, task := range tasks { + byID[task.ID] = task + } + assert.Equal(t, StatusCompleted, byID[r1.ID].Status) + assert.Equal(t, StatusCompleted, byID[r2.ID].Status) +} + +// --- Cancel Tests --- + +func TestManager_Cancel(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "cancellable", true, workBlocking()) + require.NoError(t, err) + assert.Equal(t, StatusRunning, result.Status) + + err = m.Cancel(result.ID) + require.NoError(t, err) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, StatusCanceled, task.Status) + assert.NotNil(t, task.DoneAt) + // A canceled task carries a reason rather than an empty terminal state. + assert.Equal(t, canceledError, task.Error) +} + +// A foreground run stopped by Cancel reports StatusCanceled (with the cancel +// reason) back to the caller, not StatusFailed from the work's ctx-canceled error. +func TestManager_Cancel_ForegroundReportsCanceled(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(0)}) + defer closeWithTimeout(m) + + started := make(chan string, 1) + go func() { + id := <-started + _ = m.Cancel(id) + }() + + result, err := m.Run(context.Background(), &RunInput{Description: "fg cancelable"}, + func(ctx context.Context, _ TaskInfo) (string, error) { + // Surface the task id to the canceller, then block until canceled. + for _, t := range m.List() { + started <- t.ID + } + <-ctx.Done() + return "", ctx.Err() + }) + require.NoError(t, err) + assert.Equal(t, StatusCanceled, result.Status) + assert.Equal(t, canceledError, result.Error) +} + +func TestManager_CancelNotFound(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + err := m.Cancel("nonexistent") + assert.Error(t, err) + assert.Contains(t, err.Error(), "nothing to stop") +} + +func TestManager_CancelAlreadyDone(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, _ := run(m, "task", false, workReturning("done", nil)) + + err := m.Cancel(result.ID) + assert.Error(t, err) + assert.Contains(t, err.Error(), "already finished") +} + +// --- Running-state transitions --- + +func TestManager_RunningState(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + assert.False(t, anyRunning(m)) + + result, _ := run(m, "task", true, workBlocking()) + assert.True(t, anyRunning(m)) + + _ = m.Cancel(result.ID) + waitTask(t, m, result.ID) + assert.False(t, anyRunning(m)) +} + +// --- Wait --- + +func TestManager_WaitCompleted(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "task", true, workSleeping(50*time.Millisecond, "r1")) + require.NoError(t, err) + + task := waitTask(t, m, result.ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "r1", task.Result) +} + +func TestManager_WaitTimeout(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "task", true, workBlocking()) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + task, done := m.Wait(ctx, result.ID) + require.NotNil(t, task) + assert.False(t, done) + assert.Equal(t, StatusRunning, task.Status) +} + +func TestManager_WaitNotFound(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + task, done := m.Wait(context.Background(), "missing") + assert.Nil(t, task) + assert.False(t, done) +} + +// --- Close --- + +func TestManager_Close(t *testing.T) { + m := New(context.Background(), &Config{}) + + _, _ = run(m, "task", true, workBlocking()) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + err := m.Close(ctx) + assert.NoError(t, err) + + _, err = run(m, "new", false, workReturning("x", nil)) + assert.Error(t, err) + assert.Contains(t, err.Error(), "shut down") +} + +func TestManager_RunAfterClose(t *testing.T) { + m := New(context.Background(), &Config{}) + _ = m.Close(context.Background()) + + _, err := run(m, "task", false, workReturning("x", nil)) + assert.Error(t, err) + assert.Contains(t, err.Error(), "shut down") +} + +// --- Concurrency --- + +func TestManager_ConcurrentRuns(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + const n = 50 + var wg sync.WaitGroup + wg.Add(n) + + for i := 0; i < n; i++ { + go func(i int) { + defer wg.Done() + result, err := run(m, fmt.Sprintf("task-%d", i), false, workReturning(fmt.Sprintf("result-%d", i), nil)) + require.NoError(t, err) + assert.Equal(t, StatusCompleted, result.Status) + }(i) + } + + wg.Wait() + assert.False(t, anyRunning(m)) + assert.Len(t, m.List(), n) +} + +// --- Unique IDs --- + +func TestManager_UniqueIDs(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + ids := make(map[string]bool) + for i := 0; i < 100; i++ { + result, err := run(m, "task", false, workReturning("x", nil)) + require.NoError(t, err) + assert.False(t, ids[result.ID], "duplicate ID: %s", result.ID) + ids[result.ID] = true + } +} + +// --- RunInBackground flag --- + +func TestManager_RunInBackground_Foreground(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "fg task", false, workReturning("done", nil)) + require.NoError(t, err) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.False(t, task.RunInBackground) +} + +func TestManager_RunInBackground_Background(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + result, err := run(m, "bg task", true, workSleeping(50*time.Millisecond, "bg done")) + require.NoError(t, err) + assert.Equal(t, StatusRunning, result.Status) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.True(t, task.RunInBackground) + + waitTask(t, m, result.ID) +} + +var errSentinel = errors.New("sentinel") + +func TestManager_ContextCancelStopsWork(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + started := make(chan struct{}) + work := func(ctx context.Context, _ TaskInfo) (string, error) { + close(started) + <-ctx.Done() + return "", errSentinel + } + + result, err := run(m, "task", true, work) + require.NoError(t, err) + <-started + + require.NoError(t, m.Cancel(result.ID)) + waitTask(t, m, result.ID) + + task, ok := m.Get(result.ID) + require.True(t, ok) + assert.Equal(t, StatusCanceled, task.Status) +} diff --git a/adk/backgroundtask/run_stream_test.go b/adk/backgroundtask/run_stream_test.go new file mode 100644 index 000000000..e72e229c6 --- /dev/null +++ b/adk/backgroundtask/run_stream_test.go @@ -0,0 +1,182 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package backgroundtask + +import ( + "context" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +// drainStringStream reads a string stream to EOF and returns the concatenation. +func drainStringStream(t *testing.T, sr *schema.StreamReader[string]) string { + t.Helper() + defer sr.Close() + var b strings.Builder + for { + chunk, err := sr.Recv() + if errors.Is(err, io.EOF) { + return b.String() + } + require.NoError(t, err) + b.WriteString(chunk) + } +} + +// streamWorkChunks returns a StreamWorkFunc that emits the given chunks, optionally +// pausing before each so a deadline can fire mid-stream. +func streamWorkChunks(pause time.Duration, chunks ...string) StreamWorkFunc { + return func(ctx context.Context, _ TaskInfo) (*schema.StreamReader[string], error) { + sr, sw := schema.Pipe[string](len(chunks)) + go func() { + defer sw.Close() + for _, c := range chunks { + if pause > 0 { + select { + case <-time.After(pause): + case <-ctx.Done(): + return + } + } + if sw.Send(c, nil) { + return + } + } + }() + return sr, nil + } +} + +// TestRunStream_ForegroundStreamsAndCompletes: every chunk is forwarded live and +// the accumulated text becomes the task's final Result. +func TestRunStream_ForegroundStreamsAndCompletes(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(0)}) + defer closeWithTimeout(m) + + sr, err := m.RunStream(context.Background(), &RunInput{Description: "stream"}, + streamWorkChunks(0, "a", "b", "c")) + require.NoError(t, err) + + got := drainStringStream(t, sr) + assert.Equal(t, "abc", got) + + tasks := m.List() + require.Len(t, tasks, 1) + task := waitTask(t, m, tasks[0].ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "abc", task.Result) +} + +// TestRunStream_AutoBackground: a run that outlives its budget is moved to the +// background; the caller's stream is capped with a notice and the remaining chunks +// are drained into the task Result. +func TestRunStream_AutoBackground(t *testing.T) { + m := New(context.Background(), &Config{ + ForegroundTimeoutMs: intPtr(40), + ShouldAutoBackground: func(context.Context, *Task) bool { return true }, + }) + defer closeWithTimeout(m) + + // 4 chunks, ~25ms apart; budget 40ms → ~1-2 chunks stream before background. + sr, err := m.RunStream(context.Background(), &RunInput{Description: "slow", Type: "bash"}, + streamWorkChunks(25*time.Millisecond, "1", "2", "3", "4")) + require.NoError(t, err) + + got := drainStringStream(t, sr) + assert.Contains(t, got, "moved to the background") + assert.Contains(t, got, "(bash)") + + tasks := m.List() + require.Len(t, tasks, 1) + task := waitTask(t, m, tasks[0].ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.True(t, task.RunInBackground) + // All four chunks land in the final result even though only some were streamed. + assert.Equal(t, "1234", task.Result) +} + +// TestRunStream_ExplicitBackground: no execution chunks reach the caller, only the +// notice; the work runs detached and its output becomes the task Result. +func TestRunStream_ExplicitBackground(t *testing.T) { + m := New(context.Background(), &Config{}) + defer closeWithTimeout(m) + + sr, err := m.RunStream(context.Background(), + &RunInput{Description: "bg", Type: "bash", RunInBackground: true}, + streamWorkChunks(0, "chunk-1", "chunk-2")) + require.NoError(t, err) + + got := drainStringStream(t, sr) + assert.Contains(t, got, "is running in the background") + assert.NotContains(t, got, "moved to the background") + assert.NotContains(t, got, "chunk-") + + tasks := m.List() + require.Len(t, tasks, 1) + task := waitTask(t, m, tasks[0].ID) + assert.Equal(t, StatusCompleted, task.Status) + assert.Equal(t, "chunk-1chunk-2", task.Result) +} + +// TestRunStream_WorkError: an error from the stream finalizes the task as failed +// and surfaces on the caller's stream. +func TestRunStream_WorkError(t *testing.T) { + m := New(context.Background(), &Config{ForegroundTimeoutMs: intPtr(0)}) + defer closeWithTimeout(m) + + wantErr := errors.New("boom") + work := func(ctx context.Context, _ TaskInfo) (*schema.StreamReader[string], error) { + sr, sw := schema.Pipe[string](2) + go func() { + defer sw.Close() + sw.Send("partial", nil) + sw.Send("", wantErr) + }() + return sr, nil + } + + sr, err := m.RunStream(context.Background(), &RunInput{Description: "err"}, work) + require.NoError(t, err) + + defer sr.Close() + var sawErr error + for { + _, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + if recvErr != nil { + sawErr = recvErr + break + } + } + require.Error(t, sawErr) + assert.Contains(t, sawErr.Error(), "boom") + + tasks := m.List() + require.Len(t, tasks, 1) + task := waitTask(t, m, tasks[0].ID) + assert.Equal(t, StatusFailed, task.Status) +} diff --git a/adk/call_option.go b/adk/call_option.go index 7a1cc1b65..b6f89b53c 100644 --- a/adk/call_option.go +++ b/adk/call_option.go @@ -19,12 +19,15 @@ package adk import "github.com/cloudwego/eino/callbacks" type options struct { - sharedParentSession bool - sessionValues map[string]any - checkPointID *string - skipTransferMessages bool - handlers []callbacks.Handler - cancelCtx *cancelContext + sharedParentSession bool + sessionValues map[string]any + checkPointID *string + skipTransferMessages bool + enableSessionEvents bool + enableTimelineEvents bool + enableInternalTimelineEvents bool + handlers []callbacks.Handler + cancelCtx *cancelContext } // AgentRunOption is the call option for adk Agent. @@ -55,6 +58,28 @@ func WithSessionValues(v map[string]any) AgentRunOption { }) } +func withEnableSessionEvents() AgentRunOption { + return WrapImplSpecificOptFn(func(o *options) { + o.enableSessionEvents = true + }) +} + +// WithTimelineEvents exposes the first-class SessionEvent timeline envelope on +// live AgentEvents. Without this option, lifecycle/span/observation-only events +// are still produced for managed-session persistence but are stripped from the +// user-facing stream. +func WithTimelineEvents() AgentRunOption { + return WrapImplSpecificOptFn(func(o *options) { + o.enableTimelineEvents = true + }) +} + +func withEnableInternalTimelineEvents() AgentRunOption { + return WrapImplSpecificOptFn(func(o *options) { + o.enableInternalTimelineEvents = true + }) +} + // WithSkipTransferMessages disables forwarding transfer messages during execution. // // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven diff --git a/adk/cancel_edge_test.go b/adk/cancel_edge_test.go index 0c2a80d53..1ba011cb4 100644 --- a/adk/cancel_edge_test.go +++ b/adk/cancel_edge_test.go @@ -1435,9 +1435,20 @@ func TestWithCancel_CancelImmediate_StreamableToolAborted(t *testing.T) { // ErrStreamCanceled appears on the tool's MessageStream.Recv() if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.IsStreaming && e.Output.MessageOutput.Role == schema.Tool { - // Signal that the tool stream event has been received. - close(toolStreamReady) stream := e.Output.MessageOutput.MessageStream + // Consume the first chunk so we are sure the stream is active, + // then signal readiness. This ensures cancel fires while we are + // blocked inside Recv(), preventing a race where cancel completes + // before we start consuming. + if _, firstErr := stream.Recv(); firstErr == nil { + close(toolStreamReady) + } else { + if errors.Is(firstErr, ErrStreamCanceled) { + r.foundStreamCanceled = true + } + close(toolStreamReady) + continue + } for { _, recvErr := stream.Recv() if recvErr != nil { diff --git a/adk/cancel_multicall_test.go b/adk/cancel_multicall_test.go deleted file mode 100644 index 790d14fb3..000000000 --- a/adk/cancel_multicall_test.go +++ /dev/null @@ -1,125 +0,0 @@ -/* - * Copyright 2026 CloudWeGo Authors - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package adk - -import ( - "sync/atomic" - "testing" - "time" - - "github.com/stretchr/testify/assert" - - "github.com/cloudwego/eino/compose" -) - -func TestAgentCancelFunc_MultiCall_EscalateToImmediate(t *testing.T) { - cc := newCancelContext() - var interruptCalls int32 - cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { - atomic.AddInt32(&interruptCalls, 1) - }) - cancelFn := cc.buildCancelFunc() - - handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) - handle2, _ := cancelFn(WithAgentCancelMode(CancelImmediate)) - assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls)) - - cancelErr := cc.createCancelError() - assert.Equal(t, CancelImmediate, cancelErr.Info.Mode) - assert.True(t, cancelErr.Info.Escalated) - assert.False(t, cancelErr.Info.Timeout) - - assert.True(t, cc.markCancelHandled()) - assert.NoError(t, handle1.Wait()) - assert.NoError(t, handle2.Wait()) -} - -func TestAgentCancelFunc_MultiCall_JoinSafePointModes(t *testing.T) { - cc := newCancelContext() - cancelFn := cc.buildCancelFunc() - - handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) - handle2, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) - - want := CancelAfterChatModel | CancelAfterToolCalls - assert.Equal(t, want, cc.getMode()) - - assert.True(t, cc.markCancelHandled()) - assert.NoError(t, handle1.Wait()) - assert.NoError(t, handle2.Wait()) -} - -func TestAgentCancelFunc_MultiCall_TimeoutDeadlineJoinUsesAbsoluteTime(t *testing.T) { - cc := newCancelContext() - cancelFn := cc.buildCancelFunc() - - handle1, _ := cancelFn( - WithAgentCancelMode(CancelAfterChatModel), - WithAgentCancelTimeout(200*time.Millisecond), - ) - - firstDeadline := cc.getDeadlineUnixNano() - assert.NotZero(t, firstDeadline) - - time.Sleep(50 * time.Millisecond) - - handle2, _ := cancelFn( - WithAgentCancelMode(CancelAfterToolCalls), - WithAgentCancelTimeout(60*time.Millisecond), - ) - - secondDeadline := cc.getDeadlineUnixNano() - assert.NotZero(t, secondDeadline) - assert.Less(t, secondDeadline, firstDeadline) - - assert.True(t, cc.markCancelHandled()) - assert.NoError(t, handle1.Wait()) - assert.NoError(t, handle2.Wait()) -} - -func TestAgentCancelFunc_MultiCall_TimeoutEscalationReturnsErrCancelTimeout(t *testing.T) { - cc := newCancelContext() - var interruptCalls int32 - interruptCh := make(chan struct{}, 1) - cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { - atomic.AddInt32(&interruptCalls, 1) - select { - case interruptCh <- struct{}{}: - default: - } - }) - cancelFn := cc.buildCancelFunc() - handle, _ := cancelFn( - WithAgentCancelMode(CancelAfterChatModel), - WithAgentCancelTimeout(30*time.Millisecond), - ) - - select { - case <-interruptCh: - case <-time.After(1 * time.Second): - t.Fatal("timeout escalation did not interrupt") - } - assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls)) - - cancelErr := cc.createCancelError() - assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) - assert.True(t, cancelErr.Info.Escalated) - assert.True(t, cancelErr.Info.Timeout) - - assert.True(t, cc.markCancelHandled()) - assert.Equal(t, ErrCancelTimeout, handle.Wait()) -} diff --git a/adk/cancel_recursive_test.go b/adk/cancel_recursive_test.go deleted file mode 100644 index cd7e4277f..000000000 --- a/adk/cancel_recursive_test.go +++ /dev/null @@ -1,387 +0,0 @@ -/* - * Copyright 2026 CloudWeGo Authors - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package adk - -import ( - "context" - "runtime" - "sync" - "testing" - "time" - - "github.com/stretchr/testify/assert" -) - -func assertNotClosedWithin(t *testing.T, ch <-chan struct{}, d time.Duration) { - t.Helper() - select { - case <-ch: - t.Fatal("channel was closed but should not have been") - case <-time.After(d): - } -} - -func setupParentChild(t *testing.T) (parent, child *cancelContext, cleanup func()) { - parent = newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - child = parent.deriveAgentToolCancelContext(ctx) - cleanup = func() { - child.markDone() - cancel() - } - t.Cleanup(cleanup) - return parent, child, cleanup -} - -func TestDeriveAgentToolCancelContext(t *testing.T) { - t.Run("Shallow", func(t *testing.T) { - t.Run("DoesNotPropagateSafePoint", func(t *testing.T) { - parent, child, _ := setupParentChild(t) - - parent.triggerCancel(CancelAfterChatModel) - - assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) - }) - - t.Run("ImmediateDoesNotPropagate", func(t *testing.T) { - parent, child, _ := setupParentChild(t) - - parent.triggerImmediateCancel() - - assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond) - }) - - t.Run("GrandchildNoPropagation", func(t *testing.T) { - a := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - - b := a.deriveAgentToolCancelContext(ctx) - c := b.deriveAgentToolCancelContext(ctx) - t.Cleanup(func() { - c.markDone() - b.markDone() - cancel() - }) - - a.triggerCancel(CancelAfterChatModel) - - assertNotClosedWithin(t, b.cancelChan, 50*time.Millisecond) - assertNotClosedWithin(t, c.cancelChan, 50*time.Millisecond) - }) - - t.Run("NeverRecursive_GoroutineCleanup", func(t *testing.T) { - runtime.GC() - time.Sleep(50 * time.Millisecond) - before := runtime.NumGoroutine() - - parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - - child := parent.deriveAgentToolCancelContext(ctx) - - parent.triggerCancel(CancelAfterChatModel) - time.Sleep(100 * time.Millisecond) - - child.markDone() - cancel() - - time.Sleep(200 * time.Millisecond) - runtime.GC() - time.Sleep(50 * time.Millisecond) - after := runtime.NumGoroutine() - - assert.InDelta(t, before, after, 5, "goroutine leak detected: before=%d after=%d", before, after) - }) - }) - - t.Run("Recursive", func(t *testing.T) { - t.Run("PropagatesSafePoint", func(t *testing.T) { - parent, child, _ := setupParentChild(t) - - parent.setRecursive(true) - parent.triggerCancel(CancelAfterChatModel) - - select { - case <-child.cancelChan: - case <-time.After(1 * time.Second): - t.Fatal("child did not receive cancel within 1s") - } - assert.True(t, child.shouldCancel()) - }) - - t.Run("ImmediatePropagates", func(t *testing.T) { - parent, child, _ := setupParentChild(t) - - parent.setRecursive(true) - parent.triggerImmediateCancel() - - select { - case <-child.immediateChan: - case <-time.After(1 * time.Second): - t.Fatal("child did not receive immediate cancel within 1s") - } - assert.True(t, child.isImmediateCancelled()) - }) - - t.Run("GrandchildPropagation", func(t *testing.T) { - a := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - - b := a.deriveAgentToolCancelContext(ctx) - c := b.deriveAgentToolCancelContext(ctx) - t.Cleanup(func() { - c.markDone() - b.markDone() - cancel() - }) - - a.setRecursive(true) - a.triggerCancel(CancelAfterChatModel) - - select { - case <-b.cancelChan: - case <-time.After(1 * time.Second): - t.Fatal("B did not receive cancel within 1s") - } - - select { - case <-c.cancelChan: - case <-time.After(1 * time.Second): - t.Fatal("C did not receive cancel within 1s") - } - - assert.True(t, b.shouldCancel()) - assert.True(t, c.shouldCancel()) - }) - - t.Run("SetBeforeCancel", func(t *testing.T) { - parent, child, _ := setupParentChild(t) - - parent.setRecursive(true) - - parent.triggerCancel(CancelAfterChatModel) - - select { - case <-child.cancelChan: - case <-time.After(1 * time.Second): - t.Fatal("child did not receive cancel within 1s") - } - assert.True(t, child.shouldCancel()) - }) - - t.Run("AfterRecursiveAndCancelAlreadySet", func(t *testing.T) { - parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - - parent.setRecursive(true) - parent.triggerCancel(CancelAfterChatModel) - - child := parent.deriveAgentToolCancelContext(ctx) - t.Cleanup(func() { - child.markDone() - cancel() - }) - - select { - case <-child.cancelChan: - case <-time.After(1 * time.Second): - t.Fatal("child did not immediately receive cancel") - } - assert.True(t, child.shouldCancel()) - }) - }) - - t.Run("Escalation", func(t *testing.T) { - t.Run("EscalateFromNonRecursive", func(t *testing.T) { - parent, child, _ := setupParentChild(t) - - parent.triggerCancel(CancelAfterChatModel) - - assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) - - parent.setRecursive(true) - - select { - case <-child.cancelChan: - case <-time.After(1 * time.Second): - t.Fatal("child did not receive cancel after escalation within 1s") - } - assert.True(t, child.shouldCancel()) - }) - - t.Run("EscalateImmediate", func(t *testing.T) { - parent, child, _ := setupParentChild(t) - - parent.triggerImmediateCancel() - - assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond) - - parent.setRecursive(true) - - select { - case <-child.immediateChan: - case <-time.After(1 * time.Second): - t.Fatal("child did not receive immediate cancel after escalation within 1s") - } - assert.True(t, child.isImmediateCancelled()) - }) - }) -} - -func TestDeriveAgentToolCancelContext_Race(t *testing.T) { - t.Run("SetRecursiveConcurrentWithCancelChan", func(t *testing.T) { - for i := 0; i < 100; i++ { - parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - - child := parent.deriveAgentToolCancelContext(ctx) - - var wg sync.WaitGroup - wg.Add(2) - - go func() { - defer wg.Done() - parent.setRecursive(true) - }() - - go func() { - defer wg.Done() - parent.triggerCancel(CancelAfterChatModel) - }() - - wg.Wait() - - select { - case <-child.cancelChan: - case <-time.After(1 * time.Second): - t.Fatalf("iteration %d: child did not receive cancel within 1s", i) - } - - assert.True(t, child.shouldCancel()) - child.markDone() - cancel() - } - }) - - t.Run("ChildCompletesBeforeEscalation", func(t *testing.T) { - parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - child := parent.deriveAgentToolCancelContext(ctx) - - parent.triggerCancel(CancelAfterChatModel) - time.Sleep(50 * time.Millisecond) - - child.markDone() - time.Sleep(50 * time.Millisecond) - - parent.setRecursive(true) - - assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) - }) - - t.Run("MultipleChildren_PartialCompletion", func(t *testing.T) { - parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - child1 := parent.deriveAgentToolCancelContext(ctx) - child2 := parent.deriveAgentToolCancelContext(ctx) - - parent.triggerCancel(CancelAfterChatModel) - time.Sleep(50 * time.Millisecond) - - child1.markDone() - time.Sleep(50 * time.Millisecond) - - parent.setRecursive(true) - - select { - case <-child2.cancelChan: - case <-time.After(1 * time.Second): - t.Fatal("running child did not receive cancel within 1s") - } - - assert.True(t, child2.shouldCancel()) - assert.False(t, child1.shouldCancel()) - child2.markDone() - }) - - t.Run("ContextCancelConcurrentWithRecursive", func(t *testing.T) { - done := make(chan struct{}) - go func() { - defer close(done) - - parent := newCancelContext() - ctx, cancel := context.WithCancel(context.Background()) - - child := parent.deriveAgentToolCancelContext(ctx) - - parent.triggerCancel(CancelAfterChatModel) - - var wg sync.WaitGroup - wg.Add(2) - - go func() { - defer wg.Done() - cancel() - }() - - go func() { - defer wg.Done() - parent.setRecursive(true) - }() - - wg.Wait() - child.markDone() - }() - - select { - case <-done: - case <-time.After(1 * time.Second): - t.Fatal("deadlock detected") - } - }) - - t.Run("ConcurrentSetRecursive", func(t *testing.T) { - parent := newCancelContext() - - var wg sync.WaitGroup - for i := 0; i < 10; i++ { - wg.Add(1) - go func() { - defer wg.Done() - parent.setRecursive(true) - }() - } - - done := make(chan struct{}) - go func() { - wg.Wait() - close(done) - }() - - select { - case <-done: - case <-time.After(1 * time.Second): - t.Fatal("deadlock or panic in concurrent setRecursive") - } - - assert.True(t, parent.isRecursive()) - }) -} diff --git a/adk/cancel_test.go b/adk/cancel_test.go index ea36afe11..9727db9bc 100644 --- a/adk/cancel_test.go +++ b/adk/cancel_test.go @@ -35,6 +35,62 @@ import ( "github.com/cloudwego/eino/schema" ) +type cancelAlwaysToolCallModel struct{} + +func (m *cancelAlwaysToolCallModel) Generate(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticToolCallMsg("cancel_stream_tool", "call-1", `{"input":"x"}`), nil +} + +func (m *cancelAlwaysToolCallModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.AgenticMessage{msg}), nil +} + +type cancelInterruptThenHangingStreamTool struct { + name string + interrupted chan struct{} + resumed chan struct{} + parked chan struct{} + gate chan struct{} + seen int32 + resumeOnce sync.Once + parkOnce sync.Once +} + +func (t *cancelInterruptThenHangingStreamTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "interrupt then hanging stream tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: schema.String}, + }), + }, nil +} + +func (t *cancelInterruptThenHangingStreamTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) { + if atomic.CompareAndSwapInt32(&t.seen, 0, 1) { + close(t.interrupted) + return nil, tool.StatefulInterrupt(ctx, "approval_needed", argumentsInJSON) + } + + t.resumeOnce.Do(func() { close(t.resumed) }) + r, w := schema.Pipe[string](1) + go func() { + defer w.Close() + if closed := w.Send("resumed:"+argumentsInJSON, nil); closed { + return + } + if t.parked != nil { + t.parkOnce.Do(func() { close(t.parked) }) + } + <-t.gate + }() + return r, nil +} + type cancelTestChatModel struct { delayNs int64 response *schema.Message @@ -187,6 +243,140 @@ func drainEventsAndAssertCancelError(t *testing.T, iter *AsyncIterator[*AgentEve return events } +func TestWithCancel_AgenticResumeStreamableToolTimeout_DoesNotPersistTypedNil(t *testing.T) { + ctx := context.Background() + store := newCancelTestStore() + checkpointID := "agentic-resume-streamable-tool-timeout" + streamTool := &cancelInterruptThenHangingStreamTool{ + name: "cancel_stream_tool", + interrupted: make(chan struct{}), + resumed: make(chan struct{}), + parked: make(chan struct{}), + gate: make(chan struct{}), + } + t.Cleanup(func() { + close(streamTool.gate) + }) + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "CancelAgenticResumeRepro", + Description: "repro agent", + Model: &cancelAlwaysToolCallModel{}, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{streamTool}}, + }, + }) + if err != nil { + t.Fatalf("create agent: %v", err) + } + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + CheckPointStore: store, + }) + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("go")}, WithCheckPointID(checkpointID)) + + var interruptID string + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + t.Fatalf("initial run error: %v", event.Err) + } + if event.Action == nil || event.Action.Interrupted == nil { + continue + } + for _, ictx := range event.Action.Interrupted.InterruptContexts { + if ictx.IsRootCause { + interruptID = ictx.ID + break + } + } + if interruptID != "" { + break + } + } + if interruptID == "" { + t.Fatal("root interrupt ID was not captured") + } + select { + case <-streamTool.interrupted: + case <-time.After(5 * time.Second): + t.Fatal("streamable tool did not interrupt") + } + if _, ok, getErr := store.Get(ctx, checkpointID); getErr != nil || !ok { + t.Fatalf("initial checkpoint missing: ok=%v err=%v", ok, getErr) + } + + resumeCancelOpt, resumeCancelFn := WithCancel() + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &ResumeParams{ + Targets: map[string]any{interruptID: "approved"}, + }, resumeCancelOpt) + if err != nil { + t.Fatalf("resume with params: %v", err) + } + + select { + case <-streamTool.resumed: + case <-time.After(5 * time.Second): + t.Fatal("streamable tool did not resume") + } + select { + case <-streamTool.parked: + case <-time.After(5 * time.Second): + t.Fatal("streamable tool did not park") + } + + cancelHandle, contributed := resumeCancelFn( + WithAgentCancelMode(CancelAfterToolCalls), + WithRecursive(), + WithAgentCancelTimeout(20*time.Millisecond), + ) + if !contributed { + t.Fatal("resume cancel did not contribute to active run") + } + if cancelHandle == nil { + t.Fatal("resume cancel handle is nil") + } + + cancelDone := make(chan error, 1) + go func() { + cancelDone <- cancelHandle.Wait() + }() + select { + case err = <-cancelDone: + assert.True(t, err == nil || errors.Is(err, ErrCancelTimeout) || errors.Is(err, ErrExecutionEnded), + "unexpected cancel wait error: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("resume cancel handle did not complete") + } + executionCompletedBeforeCancel := errors.Is(err, ErrExecutionEnded) + + var hasCancelError bool + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Err == nil { + continue + } + var ce *CancelError + if errors.As(event.Err, &ce) { + hasCancelError = true + } + errText := event.Err.Error() + assert.NotContains(t, errText, "gob marshal error") + assert.NotContains(t, errText, "cannot encode nil pointer") + assert.NotContains(t, errText, "*adk.agenticReactInput(nil=true") + } + assert.True(t, hasCancelError || executionCompletedBeforeCancel, + "expected CancelError in resume event stream unless execution completed before cancel") +} + func TestCancelContext(t *testing.T) { t.Run("BasicCancelContext", func(t *testing.T) { cc := newCancelContext() @@ -2297,7 +2487,7 @@ func TestCancelImmediate_OrphanedToolGoroutine_NoPanic(t *testing.T) { } assert.NotPanics(t, func() { - execCtx.send(&AgentEvent{AgentName: "test"}) + execCtx.send(context.Background(), &AgentEvent{AgentName: "test"}) }, "send after generator.Close must not panic") }) @@ -2310,21 +2500,21 @@ func TestCancelImmediate_OrphanedToolGoroutine_NoPanic(t *testing.T) { } assert.NotPanics(t, func() { - execCtx.send(&AgentEvent{AgentName: "test"}) + execCtx.send(context.Background(), &AgentEvent{AgentName: "test"}) }, "send after generator.Close must not panic even without cancelCtx (trySend safety net)") }) t.Run("unit_send_nil_execCtx", func(t *testing.T) { var execCtx *chatModelAgentExecCtx assert.NotPanics(t, func() { - execCtx.send(&AgentEvent{AgentName: "test"}) + execCtx.send(context.Background(), &AgentEvent{AgentName: "test"}) }, "send on nil execCtx must not panic") }) t.Run("unit_send_nil_generator", func(t *testing.T) { execCtx := &chatModelAgentExecCtx{} assert.NotPanics(t, func() { - execCtx.send(&AgentEvent{AgentName: "test"}) + execCtx.send(context.Background(), &AgentEvent{AgentName: "test"}) }, "send with nil generator must not panic") }) @@ -2345,7 +2535,7 @@ func TestCancelImmediate_OrphanedToolGoroutine_NoPanic(t *testing.T) { } assert.NotPanics(t, func() { - execCtx.send(&AgentEvent{AgentName: "test"}) + execCtx.send(context.Background(), &AgentEvent{AgentName: "test"}) }, "trySend must handle the case where isImmediateCancelled is false but generator is closed") }) @@ -2374,7 +2564,7 @@ func TestCancelImmediate_OrphanedToolGoroutine_NoPanic(t *testing.T) { t.Run("unit_SendEvent_no_execCtx", func(t *testing.T) { err := SendEvent(context.Background(), &AgentEvent{AgentName: "test"}) - assert.Error(t, err, "SendEvent without execCtx should return error") + assert.NoError(t, err, "SendEvent without execCtx should be a no-op") }) t.Run("integration_cancel_escalation_orphans_tool", func(t *testing.T) { @@ -3792,3 +3982,461 @@ func TestBuildCancelFunc_CASFailStateDone(t *testing.T) { t.Log("CAS race path not triggered (L743 remains a theoretical race edge)") } } + +func assertNotClosedWithin(t *testing.T, ch <-chan struct{}, d time.Duration) { + t.Helper() + select { + case <-ch: + t.Fatal("channel was closed but should not have been") + case <-time.After(d): + } +} + +func setupParentChild(t *testing.T) (parent, child *cancelContext, cleanup func()) { + parent = newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + child = parent.deriveAgentToolCancelContext(ctx) + cleanup = func() { + child.markDone() + cancel() + } + t.Cleanup(cleanup) + return parent, child, cleanup +} + +func TestDeriveAgentToolCancelContext(t *testing.T) { + t.Run("Shallow", func(t *testing.T) { + t.Run("DoesNotPropagateSafePoint", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerCancel(CancelAfterChatModel) + + assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) + }) + + t.Run("ImmediateDoesNotPropagate", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerImmediateCancel() + + assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond) + }) + + t.Run("GrandchildNoPropagation", func(t *testing.T) { + a := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + b := a.deriveAgentToolCancelContext(ctx) + c := b.deriveAgentToolCancelContext(ctx) + t.Cleanup(func() { + c.markDone() + b.markDone() + cancel() + }) + + a.triggerCancel(CancelAfterChatModel) + + assertNotClosedWithin(t, b.cancelChan, 50*time.Millisecond) + assertNotClosedWithin(t, c.cancelChan, 50*time.Millisecond) + }) + + t.Run("NeverRecursive_GoroutineCleanup", func(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + before := runtime.NumGoroutine() + + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + child := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + time.Sleep(100 * time.Millisecond) + + child.markDone() + cancel() + + time.Sleep(200 * time.Millisecond) + runtime.GC() + time.Sleep(50 * time.Millisecond) + after := runtime.NumGoroutine() + + assert.InDelta(t, before, after, 5, "goroutine leak detected: before=%d after=%d", before, after) + }) + }) + + t.Run("Recursive", func(t *testing.T) { + t.Run("PropagatesSafePoint", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.setRecursive(true) + parent.triggerCancel(CancelAfterChatModel) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive cancel within 1s") + } + assert.True(t, child.shouldCancel()) + }) + + t.Run("ImmediatePropagates", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.setRecursive(true) + parent.triggerImmediateCancel() + + select { + case <-child.immediateChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive immediate cancel within 1s") + } + assert.True(t, child.isImmediateCancelled()) + }) + + t.Run("GrandchildPropagation", func(t *testing.T) { + a := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + b := a.deriveAgentToolCancelContext(ctx) + c := b.deriveAgentToolCancelContext(ctx) + t.Cleanup(func() { + c.markDone() + b.markDone() + cancel() + }) + + a.setRecursive(true) + a.triggerCancel(CancelAfterChatModel) + + select { + case <-b.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("B did not receive cancel within 1s") + } + + select { + case <-c.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("C did not receive cancel within 1s") + } + + assert.True(t, b.shouldCancel()) + assert.True(t, c.shouldCancel()) + }) + + t.Run("SetBeforeCancel", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.setRecursive(true) + + parent.triggerCancel(CancelAfterChatModel) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive cancel within 1s") + } + assert.True(t, child.shouldCancel()) + }) + + t.Run("AfterRecursiveAndCancelAlreadySet", func(t *testing.T) { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + parent.setRecursive(true) + parent.triggerCancel(CancelAfterChatModel) + + child := parent.deriveAgentToolCancelContext(ctx) + t.Cleanup(func() { + child.markDone() + cancel() + }) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not immediately receive cancel") + } + assert.True(t, child.shouldCancel()) + }) + }) + + t.Run("Escalation", func(t *testing.T) { + t.Run("EscalateFromNonRecursive", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerCancel(CancelAfterChatModel) + + assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) + + parent.setRecursive(true) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive cancel after escalation within 1s") + } + assert.True(t, child.shouldCancel()) + }) + + t.Run("EscalateImmediate", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerImmediateCancel() + + assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond) + + parent.setRecursive(true) + + select { + case <-child.immediateChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive immediate cancel after escalation within 1s") + } + assert.True(t, child.isImmediateCancelled()) + }) + }) +} + +func TestDeriveAgentToolCancelContext_Race(t *testing.T) { + t.Run("SetRecursiveConcurrentWithCancelChan", func(t *testing.T) { + for i := 0; i < 100; i++ { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + child := parent.deriveAgentToolCancelContext(ctx) + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + parent.setRecursive(true) + }() + + go func() { + defer wg.Done() + parent.triggerCancel(CancelAfterChatModel) + }() + + wg.Wait() + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatalf("iteration %d: child did not receive cancel within 1s", i) + } + + assert.True(t, child.shouldCancel()) + child.markDone() + cancel() + } + }) + + t.Run("ChildCompletesBeforeEscalation", func(t *testing.T) { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + child := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + time.Sleep(50 * time.Millisecond) + + child.markDone() + time.Sleep(50 * time.Millisecond) + + parent.setRecursive(true) + + assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) + }) + + t.Run("MultipleChildren_PartialCompletion", func(t *testing.T) { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + child1 := parent.deriveAgentToolCancelContext(ctx) + child2 := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + time.Sleep(50 * time.Millisecond) + + child1.markDone() + time.Sleep(50 * time.Millisecond) + + parent.setRecursive(true) + + select { + case <-child2.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("running child did not receive cancel within 1s") + } + + assert.True(t, child2.shouldCancel()) + assert.False(t, child1.shouldCancel()) + child2.markDone() + }) + + t.Run("ContextCancelConcurrentWithRecursive", func(t *testing.T) { + done := make(chan struct{}) + go func() { + defer close(done) + + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + child := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + cancel() + }() + + go func() { + defer wg.Done() + parent.setRecursive(true) + }() + + wg.Wait() + child.markDone() + }() + + select { + case <-done: + case <-time.After(1 * time.Second): + t.Fatal("deadlock detected") + } + }) + + t.Run("ConcurrentSetRecursive", func(t *testing.T) { + parent := newCancelContext() + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + parent.setRecursive(true) + }() + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(1 * time.Second): + t.Fatal("deadlock or panic in concurrent setRecursive") + } + + assert.True(t, parent.isRecursive()) + }) +} + +func TestAgentCancelFunc_MultiCall_EscalateToImmediate(t *testing.T) { + cc := newCancelContext() + var interruptCalls int32 + cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { + atomic.AddInt32(&interruptCalls, 1) + }) + cancelFn := cc.buildCancelFunc() + + handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + handle2, _ := cancelFn(WithAgentCancelMode(CancelImmediate)) + assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls)) + + cancelErr := cc.createCancelError() + assert.Equal(t, CancelImmediate, cancelErr.Info.Mode) + assert.True(t, cancelErr.Info.Escalated) + assert.False(t, cancelErr.Info.Timeout) + + assert.True(t, cc.markCancelHandled()) + assert.NoError(t, handle1.Wait()) + assert.NoError(t, handle2.Wait()) +} + +func TestAgentCancelFunc_MultiCall_JoinSafePointModes(t *testing.T) { + cc := newCancelContext() + cancelFn := cc.buildCancelFunc() + + handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + handle2, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + + want := CancelAfterChatModel | CancelAfterToolCalls + assert.Equal(t, want, cc.getMode()) + + assert.True(t, cc.markCancelHandled()) + assert.NoError(t, handle1.Wait()) + assert.NoError(t, handle2.Wait()) +} + +func TestAgentCancelFunc_MultiCall_TimeoutDeadlineJoinUsesAbsoluteTime(t *testing.T) { + cc := newCancelContext() + cancelFn := cc.buildCancelFunc() + + handle1, _ := cancelFn( + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(200*time.Millisecond), + ) + + firstDeadline := cc.getDeadlineUnixNano() + assert.NotZero(t, firstDeadline) + + time.Sleep(50 * time.Millisecond) + + handle2, _ := cancelFn( + WithAgentCancelMode(CancelAfterToolCalls), + WithAgentCancelTimeout(60*time.Millisecond), + ) + + secondDeadline := cc.getDeadlineUnixNano() + assert.NotZero(t, secondDeadline) + assert.Less(t, secondDeadline, firstDeadline) + + assert.True(t, cc.markCancelHandled()) + assert.NoError(t, handle1.Wait()) + assert.NoError(t, handle2.Wait()) +} + +func TestAgentCancelFunc_MultiCall_TimeoutEscalationReturnsErrCancelTimeout(t *testing.T) { + cc := newCancelContext() + var interruptCalls int32 + interruptCh := make(chan struct{}, 1) + cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { + atomic.AddInt32(&interruptCalls, 1) + select { + case interruptCh <- struct{}{}: + default: + } + }) + cancelFn := cc.buildCancelFunc() + handle, _ := cancelFn( + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(30*time.Millisecond), + ) + + select { + case <-interruptCh: + case <-time.After(1 * time.Second): + t.Fatal("timeout escalation did not interrupt") + } + assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls)) + + cancelErr := cc.createCancelError() + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + assert.True(t, cancelErr.Info.Escalated) + assert.True(t, cancelErr.Info.Timeout) + + assert.True(t, cc.markCancelHandled()) + assert.Equal(t, ErrCancelTimeout, handle.Wait()) +} diff --git a/adk/chatmodel.go b/adk/chatmodel.go index 37a55f162..f5309eda5 100644 --- a/adk/chatmodel.go +++ b/adk/chatmodel.go @@ -19,10 +19,11 @@ package adk import ( "bytes" "context" - "encoding/gob" + "encoding/json" "errors" "fmt" "math" + "reflect" "runtime/debug" "strings" "sync" @@ -55,19 +56,112 @@ type typedChatModelAgentExecCtx[M MessageType] struct { suppressEventSend bool retryVerdictSignal *retryVerdictSignal - afterToolCallsHook func(ctx context.Context) error + afterToolCallsHook func(ctx context.Context) error + sessionEvents bool + timelineEvents bool + internalTimelineEvents bool + lastModelContext *ModelContextEvent + sawModelContext bool } -func (e *typedChatModelAgentExecCtx[M]) send(event *TypedAgentEvent[M]) { +func (e *typedChatModelAgentExecCtx[M]) send(ctx context.Context, event *TypedAgentEvent[M]) { if e == nil || e.generator == nil { return } if e.cancelCtx != nil && e.cancelCtx.isImmediateCancelled() { return } + if event == nil { + return + } + ensureTypedAgentEventMessageIDs(event) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil { + gen := sessionEventIDGeneratorFromContext[M](ctx) + if gen == nil { + gen = DefaultSessionEventIDGenerator[M] + } + if _, err := normalizeAgentSessionEventWithAssigner(event, func(se *SessionEvent[M]) (string, error) { + return gen(ctx, se) + }); err != nil { + event.Err = err + } + } e.generator.trySend(event) } +func copyModelContextEvent(event *ModelContextEvent) *ModelContextEvent { + if event == nil { + return nil + } + return &ModelContextEvent{ + ToolInfos: cloneToolInfos(event.ToolInfos), + DeferredToolInfos: cloneToolInfos(event.DeferredToolInfos), + } +} + +func cloneToolInfos(infos []*schema.ToolInfo) []*schema.ToolInfo { + if infos == nil { + return nil + } + return append([]*schema.ToolInfo{}, infos...) +} + +func modelContextEventEqual(a, b *ModelContextEvent) bool { + if a == nil || b == nil { + return a == b + } + return toolInfosEqual(a.ToolInfos, b.ToolInfos) && + toolInfosEqual(a.DeferredToolInfos, b.DeferredToolInfos) +} + +func toolInfosEqual(a, b []*schema.ToolInfo) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if !toolInfoEqual(a[i], b[i]) { + return false + } + } + return true +} + +func toolInfoEqual(a, b *schema.ToolInfo) bool { + if reflect.DeepEqual(a, b) { + return true + } + aBytes, aErr := json.Marshal(a) + bBytes, bErr := json.Marshal(b) + if aErr != nil || bErr != nil { + return false + } + return bytes.Equal(aBytes, bBytes) +} + +func syncModelContextSessionEvent[M MessageType](ctx context.Context, state *TypedChatModelAgentState[M]) { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil || !execCtx.sessionEvents || state == nil { + return + } + current := &ModelContextEvent{ + ToolInfos: cloneToolInfos(state.ToolInfos), + DeferredToolInfos: cloneToolInfos(state.DeferredToolInfos), + } + changed := !execCtx.sawModelContext || !modelContextEventEqual(execCtx.lastModelContext, current) + if changed { + execCtx.send(ctx, &TypedAgentEvent[M]{ + SessionEventVariant: &SessionEventVariant[M]{ + Event: &SessionEvent[M]{ + Kind: SessionEventModelContext, + ModelContext: copyModelContextEvent(current), + }, + }, + }) + } + execCtx.lastModelContext = copyModelContextEvent(current) + execCtx.sawModelContext = true +} + type chatModelAgentExecCtx = typedChatModelAgentExecCtx[*schema.Message] type typedChatModelAgentExecCtxKey[M MessageType] struct{} @@ -83,6 +177,43 @@ func getTypedChatModelAgentExecCtx[M MessageType](ctx context.Context) *typedCha return nil } +func newTypedChatModelAgentExecCtx[M MessageType]( + generator *AsyncGenerator[*TypedAgentEvent[M]], + cancelCtx *cancelContext, + sessionEvents bool, + timelineEvents bool, + internalTimelineEvents bool, +) *typedChatModelAgentExecCtx[M] { + return &typedChatModelAgentExecCtx[M]{ + generator: generator, + cancelCtx: cancelCtx, + sessionEvents: sessionEvents, + timelineEvents: timelineEvents, + internalTimelineEvents: internalTimelineEvents, + } +} + +func configureTypedChatModelAgentExecCtx[M MessageType]( + ctx context.Context, + generator *AsyncGenerator[*TypedAgentEvent[M]], + cancelCtx *cancelContext, + sessionEvents bool, + timelineEvents bool, + internalTimelineEvents bool, +) (context.Context, *typedChatModelAgentExecCtx[M]) { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil { + execCtx = &typedChatModelAgentExecCtx[M]{} + ctx = withTypedChatModelAgentExecCtx(ctx, execCtx) + } + execCtx.generator = generator + execCtx.cancelCtx = cancelCtx + execCtx.sessionEvents = sessionEvents + execCtx.timelineEvents = timelineEvents + execCtx.internalTimelineEvents = internalTimelineEvents + return ctx, execCtx +} + type chatModelAgentRunOptions struct { chatModelOptions []model.Option toolOptions []tool.Option @@ -90,7 +221,19 @@ type chatModelAgentRunOptions struct { historyModifier func(context.Context, []Message) []Message - afterToolCallsHook func(ctx context.Context) error + afterToolCallsHook func(ctx context.Context) error + initialModelContext *ModelContextEvent + sawInitialModelContext bool +} + +func withInitialModelContext(event *ModelContextEvent, saw bool) AgentRunOption { + return WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) { + if !saw { + return + } + t.initialModelContext = copyModelContextEvent(event) + t.sawInitialModelContext = true + }) } // WithChatModelOptions sets options for the underlying chat model. @@ -165,7 +308,7 @@ type TypedGenModelInput[M MessageType] func(ctx context.Context, instruction str type GenModelInput = TypedGenModelInput[*schema.Message] func defaultGenModelInput(ctx context.Context, instruction string, input *AgentInput) ([]Message, error) { - msgs := make([]Message, 0, len(input.Messages)+1) + inputMessages := input.Messages if instruction != "" { sp := schema.SystemMessage(instruction) @@ -184,11 +327,22 @@ func defaultGenModelInput(ctx context.Context, instruction string, input *AgentI sp = ms[0] } + // Strip any existing leading system message from history to avoid + // duplication when session state carries the previous turn's system + // message. The fresh instruction (potentially re-formatted with current + // SessionValues) always takes precedence. + if len(inputMessages) > 0 && inputMessages[0].Role == schema.System { + inputMessages = inputMessages[1:] + } + + msgs := make([]Message, 0, len(inputMessages)+1) msgs = append(msgs, sp) + msgs = append(msgs, inputMessages...) + return msgs, nil } - msgs = append(msgs, input.Messages...) - + msgs := make([]Message, 0, len(inputMessages)) + msgs = append(msgs, inputMessages...) return msgs, nil } @@ -199,11 +353,21 @@ func newDefaultGenModelInput[M MessageType]() TypedGenModelInput[M] { return any(GenModelInput(defaultGenModelInput)).(TypedGenModelInput[M]) case *schema.AgenticMessage: return any(TypedGenModelInput[*schema.AgenticMessage](func(_ context.Context, instruction string, input *TypedAgentInput[*schema.AgenticMessage]) ([]*schema.AgenticMessage, error) { - msgs := make([]*schema.AgenticMessage, 0, len(input.Messages)+1) + inputMessages := input.Messages if instruction != "" { + // Strip any existing leading system message from history to avoid + // duplication when session state carries the previous turn's system + // message. + if len(inputMessages) > 0 && inputMessages[0].Role == schema.AgenticRoleTypeSystem { + inputMessages = inputMessages[1:] + } + msgs := make([]*schema.AgenticMessage, 0, len(inputMessages)+1) msgs = append(msgs, schema.SystemAgenticMessage(instruction)) + msgs = append(msgs, inputMessages...) + return msgs, nil } - msgs = append(msgs, input.Messages...) + msgs := make([]*schema.AgenticMessage, 0, len(inputMessages)) + msgs = append(msgs, inputMessages...) return msgs, nil })).(TypedGenModelInput[M]) default: @@ -211,6 +375,78 @@ func newDefaultGenModelInput[M MessageType]() TypedGenModelInput[M] { } } +const extraKeyRuntimeGeneratedSystemMessage = "_eino_adk_runtime_generated_system_message" + +func ensureGeneratedMessageIDs[M MessageType](messages []M) { + for _, msg := range messages { + EnsureMessageID(msg) + } +} + +func isSystemRoleMessage[M MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.System + case *schema.AgenticMessage: + return m.Role == schema.AgenticRoleTypeSystem + } + return false +} + +func getMessageExtra[M MessageType](msg M) map[string]any { + switch m := any(msg).(type) { + case *schema.Message: + return m.Extra + case *schema.AgenticMessage: + return m.Extra + } + return nil +} + +func setMessageExtraField[M MessageType](msg M, key string, value any) { + switch m := any(msg).(type) { + case *schema.Message: + if m.Extra == nil { + m.Extra = map[string]any{} + } + m.Extra[key] = value + case *schema.AgenticMessage: + if m.Extra == nil { + m.Extra = map[string]any{} + } + m.Extra[key] = value + } +} + +func markRuntimeGeneratedLeadingSystem[M MessageType](inputMsgs, generatedMsgs []M) { + inputIDs := make(map[string]struct{}) + for _, msg := range inputMsgs { + if isNilMessage(msg) { + continue + } + id := GetMessageID(msg) + if id != "" { + inputIDs[id] = struct{}{} + } + } + for i := 0; i < len(generatedMsgs); i++ { + msg := generatedMsgs[i] + if isNilMessage(msg) || !isSystemRoleMessage(msg) { + break + } + id := GetMessageID(msg) + if id == "" { + setMessageExtraField(msg, extraKeyRuntimeGeneratedSystemMessage, true) + continue + } + if _, ok := inputIDs[id]; !ok { + setMessageExtraField(msg, extraKeyRuntimeGeneratedSystemMessage, true) + } else { + break + } + } +} + // TypedChatModelAgentState represents the state of a chat model agent during conversation. // This is the primary state type for both TypedChatModelAgentMiddleware and AgentMiddleware callbacks. type TypedChatModelAgentState[M MessageType] struct { @@ -357,9 +593,10 @@ type TypedChatModelAgentConfig[M MessageType] struct { // 1. eventSenderToolWrapper (internal ToolMiddleware - sends tool result events after all processing) // 2. ToolsConfig.ToolCallMiddlewares (ToolMiddleware) // 3. AgentMiddleware.WrapToolCall (ToolMiddleware) - // 4. ChatModelAgentMiddleware.WrapToolCall (wrapper, first registered is outermost) - // 5. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them) - // 6. Tool.InvokableRun/StreamableRun + // 4. cancelMonitoredToolHandler (internal - sets up cancel monitoring for stream tools) + // 5. ChatModelAgentMiddleware.WrapToolCall (wrapper, first registered is outermost) + // 6. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them) + // 7. Tool.InvokableRun/StreamableRun // // Custom Tool Event Sender Position: // By default, tool result events are emitted by an internal event sender placed before @@ -458,14 +695,17 @@ type ChatModelAgent = TypedChatModelAgent[*schema.Message] // typedRunParams holds the parameters for a typedRunFunc invocation. type typedRunParams[M MessageType] struct { - input *TypedAgentInput[M] - generator *AsyncGenerator[*TypedAgentEvent[M]] - store *bridgeStore - instruction string - returnDirectly map[string]bool - cancelCtx *cancelContext - cancelCtxOwned bool - composeOpts []compose.Option + input *TypedAgentInput[M] + generator *AsyncGenerator[*TypedAgentEvent[M]] + store *bridgeStore + instruction string + returnDirectly map[string]bool + cancelCtx *cancelContext + cancelCtxOwned bool + composeOpts []compose.Option + sessionEvents bool + timelineEvents bool + internalTimelineEvents bool afterToolCallsHook func(ctx context.Context) error } @@ -486,7 +726,7 @@ func NewChatModelAgent(ctx context.Context, config *ChatModelAgentConfig) (*Chat } // NewTypedChatModelAgent creates a new TypedChatModelAgent with the given config. -func NewTypedChatModelAgent[M MessageType](ctx context.Context, config *TypedChatModelAgentConfig[M]) (*TypedChatModelAgent[M], error) { +func NewTypedChatModelAgent[M MessageType](_ context.Context, config *TypedChatModelAgentConfig[M]) (*TypedChatModelAgent[M], error) { if config.ModelFailoverConfig != nil { if config.ModelFailoverConfig.GetFailoverModel == nil { return nil, errors.New("ModelFailoverConfig.GetFailoverModel is required when ModelFailoverConfig is set") @@ -512,21 +752,24 @@ func NewTypedChatModelAgent[M MessageType](ctx context.Context, config *TypedCha tc := config.ToolsConfig // Tool call middleware execution order (outermost to innermost): - // 1. eventSenderToolWrapper (internal - sends tool result events after all modifications) + // 1. eventSenderToolWrapper (internal — emits tool result AgentEvents and + // tool_call_start/end SessionEvents; persistence-aware so a tool call's + // start and end may straddle interrupt/resume boundaries) // 2. User-provided ToolsConfig.ToolCallMiddlewares (original order preserved) // 3. Middlewares' WrapToolCall (in registration order) - // 4. ChatModelAgentMiddleware.WrapToolCall (in registration order) - // 5. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them) + // 4. cancelMonitoredToolHandler (internal - cancel monitoring for stream tools) + // 5. ChatModelAgentMiddleware.WrapToolCall (in registration order) + // 6. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them) if !hasUserEventSenderToolWrapper(config.Handlers) { defaultToolEventSender := handlersToToolMiddlewares([]TypedChatModelAgentMiddleware[M]{newTypedEventSenderToolWrapper[M]()}) tc.ToolCallMiddlewares = append(defaultToolEventSender, tc.ToolCallMiddlewares...) } tc.ToolCallMiddlewares = append(tc.ToolCallMiddlewares, collectToolMiddlewaresFromMiddlewares(config.Middlewares)...) - // Cancel monitoring middleware (innermost — close to the tool endpoint). - // This allows early abort of the raw tool result stream when immediateChan fires - // (CancelImmediate or timeout escalation), while requiring outer wrappers to - // propagate stream errors such as ErrStreamCanceled without swallowing them. + // Cancel monitoring middleware — wraps around ChatModelAgentMiddleware handlers. + // Pre-processing sets up cancel signals; when immediateChan fires + // (CancelImmediate or timeout escalation), the raw tool result stream is aborted. + // Outer wrappers must propagate stream errors such as ErrStreamCanceled without swallowing them. cancelToolHandler := &cancelMonitoredToolHandler{} tc.ToolCallMiddlewares = append(tc.ToolCallMiddlewares, compose.ToolMiddleware{ Streamable: cancelToolHandler.WrapStreamableToolCall, @@ -790,9 +1033,12 @@ type execContext struct { toolUpdated bool // whether needs to pass a compose.WithToolList option to ToolsNode due to tool list change } -func (a *TypedChatModelAgent[M]) applyBeforeAgent(ctx context.Context, ec *execContext) (context.Context, *execContext, error) { - runCtx := &ChatModelAgentContext{ +func (a *TypedChatModelAgent[M]) applyBeforeAgent(ctx context.Context, ec *execContext, agentInput *TypedAgentInput[M]) ( + context.Context, *execContext, *TypedAgentInput[M], error) { + + runCtx := &ChatModelAgentContext[M]{ Instruction: ec.instruction, + AgentInput: agentInput, Tools: cloneSlice(ec.unwrappedTools), ReturnDirectly: copyMap(ec.returnDirectly), } @@ -801,7 +1047,7 @@ func (a *TypedChatModelAgent[M]) applyBeforeAgent(ctx context.Context, ec *execC for i, handler := range a.handlers { ctx, runCtx, err = handler.BeforeAgent(ctx, runCtx) if err != nil { - return ctx, nil, fmt.Errorf("handler[%d] (%T) BeforeAgent failed: %w", i, handler, err) + return ctx, nil, nil, fmt.Errorf("handler[%d] (%T) BeforeAgent failed: %w", i, handler, err) } } @@ -821,12 +1067,12 @@ func (a *TypedChatModelAgent[M]) applyBeforeAgent(ctx context.Context, ec *execC toolInfos, err := genToolInfos(ctx, &runtimeEC.toolsNodeConf) if err != nil { - return ctx, nil, err + return ctx, nil, nil, err } runtimeEC.toolInfos = toolInfos - return ctx, runtimeEC, nil + return ctx, runtimeEC, runCtx.AgentInput, nil } func (a *TypedChatModelAgent[M]) applyAfterAgent(ctx context.Context) (context.Context, error) { @@ -952,6 +1198,14 @@ func (a *TypedChatModelAgent[M]) handleRunFuncError( return } + if cancelCtxOwned && cancelCtx != nil && cancelCtx.shouldCancel() && errors.Is(err, ErrStreamCanceled) { + cancelErr, ok := cancelCtx.createAndMarkCancelHandled() + if ok { + generator.Send(&TypedAgentEvent[M]{Err: cancelErr}) + } + return + } + if cancelCtxOwned && cancelCtx != nil { cancelCtx.markDone() } @@ -996,6 +1250,10 @@ func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRu if err != nil { return nil, err } + if p.sessionEvents { + ensureGeneratedMessageIDs(messages) + markRuntimeGeneratedLeadingSystem(in.input.Messages, messages) + } if err := compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { st.Messages = append(st.Messages, messages...) return nil @@ -1007,18 +1265,20 @@ func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRu appendModelToChain(chain, wrappedModel) - if len(a.handlers) > 0 { - chain.AppendLambda(compose.InvokableLambda(func(ctx context.Context, msg M) (M, error) { - _, err := a.applyAfterAgent(ctx) - return msg, err - })) - } + chain.AppendLambda(compose.InvokableLambda(func(ctx context.Context, msg M) (M, error) { + if len(a.handlers) > 0 { + if _, err := a.applyAfterAgent(ctx); err != nil { + return msg, err + } + } + return msg, nil + })) var compileOptions []compose.GraphCompileOption compileOptions = append(compileOptions, compose.WithGraphName(a.name), compose.WithCheckPointStore(p.store), - compose.WithSerializer(&gobSerializer{})) + compose.WithSerializer(&schema.GobSerializer{})) if cancelCtx != nil { var interrupt func(...compose.GraphInterruptOption) @@ -1032,11 +1292,9 @@ func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRu return } - ctx = withTypedChatModelAgentExecCtx(ctx, &typedChatModelAgentExecCtx[M]{ - generator: p.generator, - cancelCtx: cancelCtx, - failoverLastSuccessModel: a.model, - }) + var execCtx *typedChatModelAgentExecCtx[M] + ctx, execCtx = configureTypedChatModelAgentExecCtx(ctx, p.generator, cancelCtx, p.sessionEvents, p.timelineEvents, p.internalTimelineEvents) + execCtx.failoverLastSuccessModel = a.model // Pre-execution cancel check if cancelCtx != nil && cancelCtx.shouldCancel() { @@ -1065,6 +1323,7 @@ func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRu err = setOutputToSession(ctx, msg, msgStream, a.outputKey) if err != nil { p.generator.Send(&TypedAgentEvent[M]{Err: err}) + return } } else if msgStream != nil { msgStream.Close() @@ -1113,13 +1372,6 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc agentName: a.name, maxIterations: a.maxIterations, } - if len(a.handlers) > 0 { - msgAgent := any(a).(*TypedChatModelAgent[*schema.Message]) - msgConf.afterAgentFunc = func(ctx context.Context, msg *schema.Message) (*schema.Message, error) { - _, err := msgAgent.applyAfterAgent(ctx) - return msg, err - } - } return func(ctx context.Context, p *typedRunParams[M]) { mp := any(p).(*typedRunParams[*schema.Message]) @@ -1130,6 +1382,17 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc } ctx = withCancelContext(ctx, cancelCtx) + msgAgent := any(a).(*TypedChatModelAgent[*schema.Message]) + msgConf.afterAgentFunc = func(ctx context.Context, msg *schema.Message) (*schema.Message, error) { + if len(a.handlers) > 0 { + _, err := msgAgent.applyAfterAgent(ctx) + if err != nil { + return msg, err + } + } + return msg, nil + } + g, err := newReact(ctx, msgConf) if err != nil { mp.generator.Send(&AgentEvent{Err: err}) @@ -1143,6 +1406,10 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc if genErr != nil { return nil, genErr } + if mp.sessionEvents { + ensureGeneratedMessageIDs(messages) + markRuntimeGeneratedLeadingSystem(in.input.Messages, messages) + } return &reactInput{ Messages: messages, }, nil @@ -1154,7 +1421,7 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc compileOptions = append(compileOptions, compose.WithGraphName(a.name), compose.WithCheckPointStore(mp.store), - compose.WithSerializer(&gobSerializer{}), + compose.WithSerializer(&schema.GobSerializer{}), compose.WithMaxRunSteps(math.MaxInt)) if cancelCtx != nil { @@ -1169,13 +1436,10 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc return } - ctx = withTypedChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{ - runtimeReturnDirectly: mp.returnDirectly, - generator: mp.generator, - cancelCtx: cancelCtx, - failoverLastSuccessModel: msgModel, - afterToolCallsHook: mp.afterToolCallsHook, - }) + ctx, execCtx := configureTypedChatModelAgentExecCtx(ctx, mp.generator, cancelCtx, mp.sessionEvents, mp.timelineEvents, mp.internalTimelineEvents) + execCtx.runtimeReturnDirectly = mp.returnDirectly + execCtx.failoverLastSuccessModel = msgModel + execCtx.afterToolCallsHook = mp.afterToolCallsHook // Pre-execution cancel check if cancelCtx != nil && cancelCtx.shouldCancel() { @@ -1216,6 +1480,7 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc err_ = setOutputToSession[*schema.Message](ctx, msg, msgStream, a.outputKey) if err_ != nil { mp.generator.Send(&AgentEvent{Err: err_}) + return } } else if msgStream != nil { msgStream.Close() @@ -1251,13 +1516,6 @@ func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(_ context.Context, bc agentName: a.name, maxIterations: a.maxIterations, } - if len(a.handlers) > 0 { - agenticAgent := any(a).(*TypedChatModelAgent[*schema.AgenticMessage]) - agenticConf.afterAgentFunc = func(ctx context.Context, msg *schema.AgenticMessage) (*schema.AgenticMessage, error) { - _, err := agenticAgent.applyAfterAgent(ctx) - return msg, err - } - } return func(ctx context.Context, p *typedRunParams[M]) { ap := any(p).(*typedRunParams[*schema.AgenticMessage]) @@ -1268,6 +1526,17 @@ func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(_ context.Context, bc } ctx = withCancelContext(ctx, cancelCtx) + agenticAgent := any(a).(*TypedChatModelAgent[*schema.AgenticMessage]) + agenticConf.afterAgentFunc = func(ctx context.Context, msg *schema.AgenticMessage) (*schema.AgenticMessage, error) { + if len(a.handlers) > 0 { + _, err := agenticAgent.applyAfterAgent(ctx) + if err != nil { + return msg, err + } + } + return msg, nil + } + g, err := newAgenticReact(ctx, agenticConf) if err != nil { ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err}) @@ -1281,18 +1550,22 @@ func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(_ context.Context, bc if genErr != nil { return nil, genErr } + if ap.sessionEvents { + ensureGeneratedMessageIDs(messages) + markRuntimeGeneratedLeadingSystem(in.input.Messages, messages) + } return &agenticReactInput{ Messages: messages, }, nil }), ). - AppendGraph(g, compose.WithNodeName("ReAct"), compose.WithGraphCompileOptions(compose.WithMaxRunSteps(math.MaxInt))) + AppendGraph(g, compose.WithNodeName("AgenticReAct"), compose.WithGraphCompileOptions(compose.WithMaxRunSteps(math.MaxInt))) var compileOptions []compose.GraphCompileOption compileOptions = append(compileOptions, compose.WithGraphName(a.name), compose.WithCheckPointStore(ap.store), - compose.WithSerializer(&gobSerializer{}), + compose.WithSerializer(&schema.GobSerializer{}), compose.WithMaxRunSteps(math.MaxInt)) if cancelCtx != nil { @@ -1307,13 +1580,10 @@ func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(_ context.Context, bc return } - ctx = withTypedChatModelAgentExecCtx(ctx, &typedChatModelAgentExecCtx[*schema.AgenticMessage]{ - runtimeReturnDirectly: ap.returnDirectly, - generator: ap.generator, - cancelCtx: cancelCtx, - failoverLastSuccessModel: agenticModel, - afterToolCallsHook: ap.afterToolCallsHook, - }) + ctx, execCtx := configureTypedChatModelAgentExecCtx(ctx, ap.generator, cancelCtx, ap.sessionEvents, ap.timelineEvents, ap.internalTimelineEvents) + execCtx.runtimeReturnDirectly = ap.returnDirectly + execCtx.failoverLastSuccessModel = agenticModel + execCtx.afterToolCallsHook = ap.afterToolCallsHook // Pre-execution cancel check if cancelCtx != nil && cancelCtx.shouldCancel() { @@ -1351,6 +1621,7 @@ func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(_ context.Context, bc err_ = setOutputToSession(ctx, msg, msgStream, a.outputKey) if err_ != nil { ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err_}) + return } } else if msgStream != nil { msgStream.Close() @@ -1398,12 +1669,12 @@ func (a *TypedChatModelAgent[M]) buildRunFunc(ctx context.Context) typedRunFunc[ return a.run } -func (a *TypedChatModelAgent[M]) getRunFunc(ctx context.Context) (context.Context, typedRunFunc[M], *execContext, error) { +func (a *TypedChatModelAgent[M]) getRunFunc(ctx context.Context, agentInput *TypedAgentInput[M]) (context.Context, typedRunFunc[M], *execContext, *TypedAgentInput[M], error) { defaultRun := a.buildRunFunc(ctx) bc := a.exeCtx if bc == nil { - return ctx, defaultRun, bc, nil + return ctx, defaultRun, bc, agentInput, nil } if len(a.handlers) == 0 { @@ -1413,32 +1684,32 @@ func (a *TypedChatModelAgent[M]) getRunFunc(ctx context.Context) (context.Contex returnDirectly: bc.returnDirectly, toolInfos: bc.toolInfos, } - return ctx, defaultRun, runtimeBC, nil + return ctx, defaultRun, runtimeBC, agentInput, nil } - ctx, runtimeBC, err := a.applyBeforeAgent(ctx, bc) + ctx, runtimeBC, agentInput, err := a.applyBeforeAgent(ctx, bc, agentInput) if err != nil { - return ctx, nil, nil, err + return ctx, nil, nil, nil, err } if !runtimeBC.rebuildGraph { - return ctx, defaultRun, runtimeBC, nil + return ctx, defaultRun, runtimeBC, agentInput, nil } var tempRun typedRunFunc[M] if len(runtimeBC.toolsNodeConf.Tools) == 0 { tempRun, err = a.buildNoToolsRunFunc(ctx) if err != nil { - return ctx, nil, nil, err + return ctx, nil, nil, nil, err } } else { tempRun, err = a.buildReActRunFunc(ctx, runtimeBC) if err != nil { - return ctx, nil, nil, err + return ctx, nil, nil, nil, err } } - return ctx, tempRun, runtimeBC, nil + return ctx, tempRun, runtimeBC, agentInput, nil } func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput[M], opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { @@ -1446,8 +1717,10 @@ func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput o := getCommonOptions(nil, opts...) cancelCtx, cancelCtxOwned := resolveRunCancelContext(ctx, o) + ctx = withTypedChatModelAgentExecCtx(ctx, + newTypedChatModelAgentExecCtx(generator, cancelCtx, o.enableSessionEvents, o.enableTimelineEvents, o.enableInternalTimelineEvents)) - ctx, run, bc, err := a.getRunFunc(ctx) + ctx, run, bc, input, err := a.getRunFunc(ctx, input) if err != nil { go func() { if cancelCtxOwned && cancelCtx != nil { @@ -1462,6 +1735,10 @@ func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput co := getComposeOptions(opts) co = append(co, compose.WithCheckPointID(bridgeCheckpointID)) runOps := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...) + if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil { + execCtx.lastModelContext = copyModelContextEvent(runOps.initialModelContext) + execCtx.sawModelContext = runOps.sawInitialModelContext + } if bc != nil { if len(bc.toolInfos) > 0 { @@ -1497,15 +1774,18 @@ func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput } run(ctx, &typedRunParams[M]{ - input: input, - generator: generator, - store: newBridgeStore(), - instruction: instruction, - returnDirectly: returnDirectly, - cancelCtx: cancelCtx, - cancelCtxOwned: cancelCtxOwned, - composeOpts: co, - afterToolCallsHook: runOps.afterToolCallsHook, + input: input, + generator: generator, + store: newBridgeStore(), + instruction: instruction, + returnDirectly: returnDirectly, + cancelCtx: cancelCtx, + cancelCtxOwned: cancelCtxOwned, + composeOpts: co, + sessionEvents: o.enableSessionEvents, + timelineEvents: o.enableTimelineEvents, + internalTimelineEvents: o.enableInternalTimelineEvents, + afterToolCallsHook: runOps.afterToolCallsHook, }) }() @@ -1520,8 +1800,10 @@ func (a *TypedChatModelAgent[M]) Resume(ctx context.Context, info *ResumeInfo, o o := getCommonOptions(nil, opts...) cancelCtx, cancelCtxOwned := resolveRunCancelContext(ctx, o) + ctx = withTypedChatModelAgentExecCtx(ctx, + newTypedChatModelAgentExecCtx(generator, cancelCtx, o.enableSessionEvents, o.enableTimelineEvents, o.enableInternalTimelineEvents)) - ctx, run, bc, err := a.getRunFunc(ctx) + ctx, run, bc, _, err := a.getRunFunc(ctx, nil) if err != nil { go func() { if cancelCtxOwned && cancelCtx != nil { @@ -1536,6 +1818,10 @@ func (a *TypedChatModelAgent[M]) Resume(ctx context.Context, info *ResumeInfo, o co := getComposeOptions(opts) co = append(co, compose.WithCheckPointID(bridgeCheckpointID)) resumeRunOps := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...) + if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil { + execCtx.lastModelContext = copyModelContextEvent(resumeRunOps.initialModelContext) + execCtx.sawModelContext = resumeRunOps.sawInitialModelContext + } if bc != nil { if len(bc.toolInfos) > 0 { @@ -1621,15 +1907,18 @@ func (a *TypedChatModelAgent[M]) Resume(ctx context.Context, info *ResumeInfo, o } run(ctx, &typedRunParams[M]{ - input: &TypedAgentInput[M]{EnableStreaming: info.EnableStreaming}, - generator: generator, - store: newResumeBridgeStore(bridgeCheckpointID, stateByte), - instruction: instruction, - returnDirectly: returnDirectly, - cancelCtx: cancelCtx, - cancelCtxOwned: cancelCtxOwned, - composeOpts: co, - afterToolCallsHook: resumeRunOps.afterToolCallsHook, + input: &TypedAgentInput[M]{EnableStreaming: info.EnableStreaming}, + generator: generator, + store: newResumeBridgeStore(bridgeCheckpointID, stateByte), + instruction: instruction, + returnDirectly: returnDirectly, + cancelCtx: cancelCtx, + cancelCtxOwned: cancelCtxOwned, + composeOpts: co, + sessionEvents: o.enableSessionEvents, + timelineEvents: o.enableTimelineEvents, + internalTimelineEvents: o.enableInternalTimelineEvents, + afterToolCallsHook: resumeRunOps.afterToolCallsHook, }) }() @@ -1668,22 +1957,6 @@ func getComposeOptions(opts []AgentRunOption) []compose.Option { return co } -type gobSerializer struct{} - -func (g *gobSerializer) Marshal(v any) ([]byte, error) { - buf := new(bytes.Buffer) - err := gob.NewEncoder(buf).Encode(v) - if err != nil { - return nil, err - } - return buf.Bytes(), nil -} - -func (g *gobSerializer) Unmarshal(data []byte, v any) error { - buf := bytes.NewBuffer(data) - return gob.NewDecoder(buf).Decode(v) -} - // preprocessComposeCheckpoint migrates legacy compose checkpoints to the current format. // It handles the v0.8.0-v0.8.3 format: // - gob name "_eino_adk_state_v080_" (already byte-patched by preprocessADKCheckpoint @@ -1697,7 +1970,7 @@ func preprocessComposeCheckpoint(data []byte) ([]byte, error) { const lenPrefixedCompatName = "\x15" + stateGobNameV080 if bytes.Contains(data, []byte(lenPrefixedCompatName)) { // v0.8.0-v0.8.3: already byte-patched by preprocessADKCheckpoint; decode as *stateV080. - migrated, err := compose.MigrateCheckpointState(data, &gobSerializer{}, func(state any) (any, bool, error) { + migrated, err := compose.MigrateCheckpointState(data, &schema.GobSerializer{}, func(state any) (any, bool, error) { sc, ok := state.(*stateV080) if !ok { return state, false, nil diff --git a/adk/chatmodel_retry_test.go b/adk/chatmodel_retry_test.go index e7ef1592f..f6ada0d10 100644 --- a/adk/chatmodel_retry_test.go +++ b/adk/chatmodel_retry_test.go @@ -2602,7 +2602,7 @@ func TestErrStreamCanceled(t *testing.T) { }) } -func TestAttack_ShouldRetry_NilDecisionOnEveryCall(t *testing.T) { +func TestRetryChatModel_ShouldRetryNilDecisionOnEveryCall(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -2645,7 +2645,7 @@ func TestAttack_ShouldRetry_NilDecisionOnEveryCall(t *testing.T) { assert.True(t, foundOK, "nil decision should accept the message as-is") } -func TestAttack_ShouldRetry_MaxRetriesZero_RejectFirstAttempt(t *testing.T) { +func TestRetryChatModel_ShouldRetryMaxRetriesZeroRejectFirstAttempt(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -2685,7 +2685,7 @@ func TestAttack_ShouldRetry_MaxRetriesZero_RejectFirstAttempt(t *testing.T) { assert.True(t, foundExhausted, "MaxRetries=0 with Retry:true should produce RetryExhaustedError") } -func TestAttack_ShouldRetry_RetryTrueWithRewriteError_IgnoresRewrite(t *testing.T) { +func TestRetryChatModel_ShouldRetryTrueWithRewriteErrorIgnoresRewrite(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -2738,7 +2738,7 @@ func TestAttack_ShouldRetry_RetryTrueWithRewriteError_IgnoresRewrite(t *testing. assert.True(t, foundSuccess, "should eventually succeed after retry, ignoring RewriteError") } -func TestAttack_ShouldRetry_OptionsAccumulateAcrossRetries(t *testing.T) { +func TestRetryChatModel_ShouldRetryOptionsAccumulateAcrossRetries(t *testing.T) { ctx := context.Background() var capturedOpts [][]model.Option @@ -2789,7 +2789,7 @@ func TestAttack_ShouldRetry_OptionsAccumulateAcrossRetries(t *testing.T) { "third call should have more options than second (accumulated AdditionalOptions)") } -func TestAttack_ShouldRetry_Stream_NilDecisionAccepts(t *testing.T) { +func TestRetryChatModel_ShouldRetryStreamNilDecisionAccepts(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -2835,7 +2835,7 @@ func TestAttack_ShouldRetry_Stream_NilDecisionAccepts(t *testing.T) { } } -func TestAttack_ShouldRetry_Stream_MaxRetriesZero_Exhausted(t *testing.T) { +func TestRetryChatModel_ShouldRetryStreamMaxRetriesZeroExhausted(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -2893,7 +2893,7 @@ func TestAttack_ShouldRetry_Stream_MaxRetriesZero_Exhausted(t *testing.T) { assert.True(t, foundExhausted, "MaxRetries=0 stream reject should produce RetryExhaustedError") } -func TestAttack_ShouldRetry_Stream_RewriteErrorOnCleanStream(t *testing.T) { +func TestRetryChatModel_ShouldRetryStreamRewriteErrorOnCleanStream(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -2949,7 +2949,7 @@ func TestAttack_ShouldRetry_Stream_RewriteErrorOnCleanStream(t *testing.T) { assert.True(t, foundFatal, "RewriteError on clean stream should propagate the fatal error") } -func TestAttack_ShouldRetry_ConcatMessagesFails_EmptyStream(t *testing.T) { +func TestRetryChatModel_ShouldRetryConcatMessagesFailsEmptyStream(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -3003,7 +3003,7 @@ func TestAttack_ShouldRetry_ConcatMessagesFails_EmptyStream(t *testing.T) { assert.Nil(t, capturedCtx.Err, "empty stream should have nil Err") } -func TestAttack_ShouldRetry_Stream_MidStreamError_VerdictDoubleRead(t *testing.T) { +func TestRetryChatModel_ShouldRetryStreamMidStreamErrorVerdictDoubleRead(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() diff --git a/adk/chatmodel_test.go b/adk/chatmodel_test.go index 2c9206478..6e295d751 100644 --- a/adk/chatmodel_test.go +++ b/adk/chatmodel_test.go @@ -86,6 +86,48 @@ func TestChatModelAgentRun(t *testing.T) { assert.False(t, ok) }) + t.Run("SessionEvents_NoTools_EmitsModelContext", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("session answer", nil), nil). + Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "SessionAgent", + Description: "session event test agent", + Instruction: "You are a helpful assistant.", + Model: cm, + OutputKey: "answer", + }) + require.NoError(t, err) + + input := &AgentInput{Messages: []Message{schema.UserMessage("remember this")}} + ctx = ctxWithNewTypedRunCtx(ctx, input, false) + + iterator := agent.Run(ctx, input, withEnableSessionEvents()) + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + require.NoError(t, event.Err) + events = append(events, event) + } + + require.Len(t, events, 2) + require.NotNil(t, events[0].SessionEventVariant.Event) + assert.Equal(t, SessionEventModelContext, events[0].SessionEventVariant.Event.Kind) + require.NotNil(t, events[0].SessionEventVariant.Event.ModelContext) + assert.Empty(t, events[0].SessionEventVariant.Event.ModelContext.ToolInfos) + + require.NotNil(t, events[1].Output) + assert.Equal(t, "session answer", events[1].Output.MessageOutput.Message.Content) + }) + t.Run("BasicChatModelWithAgentMiddleware", func(t *testing.T) { ctx := context.Background() @@ -204,6 +246,62 @@ func TestChatModelAgentRun(t *testing.T) { assert.Len(t, capturedMessages, 3) }) + t.Run("SessionEvents_ReAct_EmitsToolAwareModelContext", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("need tool", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: "test_tool", Arguments: "{}"}}, + }), nil + } + return schema.AssistantMessage("final with tool", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "SessionReActAgent", + Description: "session react event test agent", + Instruction: "You are a helpful assistant.", + Model: cm, + OutputKey: "answer", + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&fakeToolForTest{tarCount: 0}}, + }, + }, + }) + require.NoError(t, err) + + input := &AgentInput{Messages: []Message{schema.UserMessage("use tool")}} + ctx = ctxWithNewTypedRunCtx(ctx, input, false) + + iterator := agent.Run(ctx, input, withEnableSessionEvents()) + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + require.NoError(t, event.Err) + events = append(events, event) + } + + require.Len(t, events, 4) + assert.Equal(t, 2, generateCount) + require.NotNil(t, events[0].SessionEventVariant.Event) + assert.Equal(t, SessionEventModelContext, events[0].SessionEventVariant.Event.Kind) + require.NotNil(t, events[0].SessionEventVariant.Event.ModelContext) + require.Len(t, events[0].SessionEventVariant.Event.ModelContext.ToolInfos, 1) + assert.Equal(t, "test_tool", events[0].SessionEventVariant.Event.ModelContext.ToolInfos[0].Name) + }) + t.Run("AfterChatModel_ReAct_ModifyAffectsFlow", func(t *testing.T) { ctx := context.Background() @@ -2415,3 +2513,122 @@ func TestToolAliasesPropagation(t *testing.T) { assert.NotContains(t, args, "grep_content") }) } + +func TestMarkRuntimeGeneratedLeadingSystem(t *testing.T) { + t.Run("no system messages in generated output", func(t *testing.T) { + input := []*schema.Message{schema.UserMessage("hi")} + generated := []*schema.Message{schema.UserMessage("hi"), schema.AssistantMessage("hello", nil)} + markRuntimeGeneratedLeadingSystem(input, generated) + for _, msg := range generated { + _, ok := msg.Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.False(t, ok, "no system messages should be marked") + } + }) + + t.Run("single generated system message with empty input", func(t *testing.T) { + input := []*schema.Message{} + sys := schema.SystemMessage("gen-sys") + EnsureMessageID(sys) + generated := []*schema.Message{sys, schema.UserMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + val, ok := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.True(t, ok, "generated system message should be marked") + assert.Equal(t, true, val) + _, ok = generated[1].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.False(t, ok, "non-system message should not be marked") + }) + + t.Run("multiple generated system messages", func(t *testing.T) { + input := []*schema.Message{} + sys1 := schema.SystemMessage("sys1") + EnsureMessageID(sys1) + sys2 := schema.SystemMessage("sys2") + EnsureMessageID(sys2) + generated := []*schema.Message{sys1, sys2, schema.UserMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + val1, ok1 := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.True(t, ok1) + assert.Equal(t, true, val1) + val2, ok2 := generated[1].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.True(t, ok2) + assert.Equal(t, true, val2) + }) + + t.Run("caller-supplied system message passed through unchanged", func(t *testing.T) { + callerSys := schema.SystemMessage("caller-sys") + EnsureMessageID(callerSys) + input := []*schema.Message{callerSys, schema.UserMessage("hi")} + generated := []*schema.Message{callerSys, schema.UserMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + _, ok := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.False(t, ok, "caller-supplied system message should not be marked") + }) + + t.Run("new system prepended before caller-supplied system", func(t *testing.T) { + callerSys := schema.SystemMessage("caller-sys") + EnsureMessageID(callerSys) + input := []*schema.Message{callerSys, schema.UserMessage("hi")} + newSys := schema.SystemMessage("new-sys") + EnsureMessageID(newSys) + generated := []*schema.Message{newSys, callerSys, schema.UserMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + val, ok := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.True(t, ok, "newly prepended system message should be marked") + assert.Equal(t, true, val) + _, ok = generated[1].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.False(t, ok, "caller-supplied passthrough system should not be marked") + }) + + t.Run("input has system but generated has different one", func(t *testing.T) { + callerSys := schema.SystemMessage("old") + EnsureMessageID(callerSys) + input := []*schema.Message{callerSys, schema.UserMessage("hi")} + newSys := schema.SystemMessage("new") + EnsureMessageID(newSys) + generated := []*schema.Message{newSys, schema.UserMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + val, ok := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.True(t, ok, "different system message should be marked as runtime-generated") + assert.Equal(t, true, val) + }) + + t.Run("agentic messages - generated system marked", func(t *testing.T) { + input := []*schema.AgenticMessage{} + sys := schema.SystemAgenticMessage("gen-sys") + EnsureMessageID(sys) + generated := []*schema.AgenticMessage{sys, schema.UserAgenticMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + val, ok := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.True(t, ok) + assert.Equal(t, true, val) + }) + + t.Run("agentic messages - passthrough not marked", func(t *testing.T) { + callerSys := schema.SystemAgenticMessage("caller-sys") + EnsureMessageID(callerSys) + input := []*schema.AgenticMessage{callerSys} + generated := []*schema.AgenticMessage{callerSys, schema.UserAgenticMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + _, ok := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.False(t, ok) + }) + + t.Run("generated messages have no IDs - all leading systems marked", func(t *testing.T) { + callerSys := schema.SystemMessage("caller") + EnsureMessageID(callerSys) + input := []*schema.Message{callerSys} + genSys := schema.SystemMessage("gen") + generated := []*schema.Message{genSys, schema.UserMessage("hi")} + markRuntimeGeneratedLeadingSystem(input, generated) + val, ok := generated[0].Extra[extraKeyRuntimeGeneratedSystemMessage] + assert.True(t, ok, "system with no ID should be marked (no durable identity match)") + assert.Equal(t, true, val) + }) + + t.Run("empty generated messages", func(t *testing.T) { + input := []*schema.Message{schema.SystemMessage("sys")} + generated := []*schema.Message{} + markRuntimeGeneratedLeadingSystem(input, generated) + assert.Empty(t, generated) + }) +} diff --git a/adk/config.go b/adk/config.go index 67ead86bb..6fb2ec53b 100644 --- a/adk/config.go +++ b/adk/config.go @@ -21,6 +21,9 @@ import "github.com/cloudwego/eino/adk/internal" // Language represents the language setting for the ADK built-in prompts. type Language = internal.Language +// I18nPrompts holds prompt strings for different languages. +type I18nPrompts = internal.I18nPrompts + const ( // LanguageEnglish represents English language. LanguageEnglish Language = internal.LanguageEnglish @@ -33,3 +36,8 @@ const ( func SetLanguage(lang Language) error { return internal.SetLanguage(lang) } + +// SelectPrompt returns the prompt string for the current ADK built-in prompt language. +func SelectPrompt(prompts I18nPrompts) string { + return internal.SelectPrompt(prompts) +} diff --git a/adk/config_test.go b/adk/config_test.go new file mode 100644 index 000000000..5570b1b01 --- /dev/null +++ b/adk/config_test.go @@ -0,0 +1,45 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPromptLanguageWrappers(t *testing.T) { + require.NoError(t, SetLanguage(LanguageEnglish)) + t.Cleanup(func() { + require.NoError(t, SetLanguage(LanguageEnglish)) + }) + + prompts := I18nPrompts{ + English: "hello", + Chinese: "你好", + } + + assert.Equal(t, "hello", SelectPrompt(prompts)) + + require.NoError(t, SetLanguage(LanguageChinese)) + assert.Equal(t, "你好", SelectPrompt(prompts)) + + err := SetLanguage(Language(255)) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid language") +} diff --git a/adk/coverage_contract_test.go b/adk/coverage_contract_test.go new file mode 100644 index 000000000..0cba9d63f --- /dev/null +++ b/adk/coverage_contract_test.go @@ -0,0 +1,218 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" + "github.com/cloudwego/eino/schema/claude" + "github.com/cloudwego/eino/schema/gemini" +) + +type serviceContractStore struct { + loadSessionIDs []string + loadReqs []*LoadSessionEventsRequest + appendSessionIDs []string + appendEvents [][]*SessionEvent[*schema.Message] + loadErr error + appendErr error +} + +func (s *serviceContractStore) LoadEvents(_ context.Context, sessionID string, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + s.loadSessionIDs = append(s.loadSessionIDs, sessionID) + s.loadReqs = append(s.loadReqs, req) + if s.loadErr != nil { + return nil, s.loadErr + } + return &LoadSessionEventsResult[*schema.Message]{}, nil +} + +func (s *serviceContractStore) AppendEvents(_ context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + s.appendSessionIDs = append(s.appendSessionIDs, sessionID) + s.appendEvents = append(s.appendEvents, events) + if s.appendErr != nil { + return s.appendErr + } + return nil +} + +func TestUsageHelpersExtractAssistantMetadata(t *testing.T) { + usage := &schema.TokenUsage{ + PromptTokens: 11, + CompletionTokens: 7, + PromptTokenDetails: schema.PromptTokenDetails{ + CachedTokens: 5, + }, + } + msg := schema.AssistantMessage("ok", nil) + msg.ResponseMeta = &schema.ResponseMeta{FinishReason: "stop", Usage: usage} + + assert.Same(t, usage, assistantTokenUsage[*schema.Message](msg)) + assert.Equal(t, "stop", assistantFinishReason[*schema.Message](msg)) + got := modelUsageFromAssistant[*schema.Message](msg) + require.NotNil(t, got) + assert.Equal(t, 11, got.InputTokens) + assert.Equal(t, 7, got.OutputTokens) + assert.Equal(t, 5, got.CacheReadInputTokens) + assert.Same(t, usage, got.Raw) + + assert.Nil(t, assistantTokenUsage[*schema.Message](schema.UserMessage("q"))) + assert.Empty(t, assistantFinishReason[*schema.Message](schema.UserMessage("q"))) + assert.Nil(t, modelUsageFromAssistant[*schema.Message](schema.AssistantMessage("no usage", nil))) + assert.Nil(t, assistantTokenUsage[*schema.Message](nil)) + + agenticUsage := &schema.TokenUsage{PromptTokens: 3, CompletionTokens: 4} + agentic := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ResponseMeta: &schema.AgenticResponseMeta{ + TokenUsage: agenticUsage, + ClaudeExtension: &claude.ResponseMetaExtension{StopReason: "end_turn"}, + }, + } + assert.Same(t, agenticUsage, assistantTokenUsage[*schema.AgenticMessage](agentic)) + assert.Equal(t, "end_turn", assistantFinishReason[*schema.AgenticMessage](agentic)) + + agentic.ResponseMeta.ClaudeExtension = nil + agentic.ResponseMeta.GeminiExtension = &gemini.ResponseMetaExtension{FinishReason: "STOP"} + assert.Equal(t, "STOP", assistantFinishReason[*schema.AgenticMessage](agentic)) + assert.Empty(t, assistantFinishReason[*schema.AgenticMessage](&schema.AgenticMessage{Role: schema.AgenticRoleTypeUser})) + assert.Nil(t, assistantTokenUsage[*schema.AgenticMessage](nil)) +} + +func TestCommonOptionsAndFilteringContracts(t *testing.T) { + values := map[string]any{"k": "v"} + base := getCommonOptions(nil, + WithSessionValues(values), + withEnableSessionEvents(), + WithTimelineEvents(), + withEnableInternalTimelineEvents(), + WithSkipTransferMessages(), + withSharedParentSession(), + WithCallbacks(nil), + ) + require.NotNil(t, base) + assert.Equal(t, values, base.sessionValues) + assert.True(t, base.enableSessionEvents) + assert.True(t, base.enableTimelineEvents) + assert.True(t, base.enableInternalTimelineEvents) + assert.True(t, base.skipTransferMessages) + assert.True(t, base.sharedParentSession) + assert.Len(t, base.handlers, 1) + + custom := GetImplSpecificOptions(&struct{ Seen bool }{}, WrapImplSpecificOptFn(func(o *struct{ Seen bool }) { + o.Seen = true + })) + assert.True(t, custom.Seen) + + undesignatedCallback := WithCallbacks(nil) + currentCallback := WithCallbacks(nil).DesignateAgent("parent") + otherCallback := WithCallbacks(nil).DesignateAgent("child") + nonCallback := WithSessionValues(map[string]any{"x": "y"}) + filtered := filterCallbackHandlersForNestedAgents("parent", []AgentRunOption{ + undesignatedCallback, + currentCallback, + otherCallback, + nonCallback, + {}, + }) + assert.Len(t, filtered, 3) + assert.Equal(t, []string{"child"}, filtered[0].agentNames) + assert.NotNil(t, filtered[1].implSpecificOptFn) + assert.Nil(t, filtered[2].implSpecificOptFn) + + cancelOpt := WrapImplSpecificOptFn(func(o *options) { + o.cancelCtx = &cancelContext{} + }) + filtered = filterCancelOption([]AgentRunOption{cancelOpt, nonCallback, {}}) + assert.Len(t, filtered, 2) + assert.NotNil(t, filtered[0].implSpecificOptFn) + assert.Nil(t, filtered[1].implSpecificOptFn) + + assert.Nil(t, filterCallbackHandlersForNestedAgents("parent", nil)) + assert.Nil(t, filterCancelOption(nil)) + assert.Nil(t, filterOptions("parent", nil)) + assert.Len(t, filterOptions("parent", []AgentRunOption{nonCallback.DesignateAgent("parent"), otherCallback, {}}), 2) +} + +func TestLocalSessionStoreHandleContracts(t *testing.T) { + ctx := context.Background() + var nilStore SessionEventStore[*schema.Message] + _, err := openLocalSession[*schema.Message](ctx, nilStore, &openSessionRequest{sessionID: "sid"}) + require.ErrorIs(t, err, ErrSessionBusy) + + store := &serviceContractStore{} + require.NotNil(t, store) + + _, err = openLocalSession[*schema.Message](ctx, store, nil) + require.ErrorIs(t, err, ErrSessionBusy) + _, err = openLocalSession[*schema.Message](ctx, store, &openSessionRequest{}) + require.ErrorIs(t, err, ErrSessionBusy) + + opened, err := openLocalSession[*schema.Message](ctx, store, &openSessionRequest{sessionID: "sid"}) + require.NoError(t, err) + require.NotNil(t, opened) + + _, err = openLocalSession[*schema.Message](ctx, store, &openSessionRequest{sessionID: "sid"}) + require.ErrorIs(t, err, ErrSessionBusy) + + res, err := opened.handle.loadEvents(ctx, nil) + require.NoError(t, err) + require.NotNil(t, res) + require.Len(t, store.loadReqs, 1) + assert.Equal(t, "sid", store.loadSessionIDs[0]) + + event := validTestPayload() + err = opened.handle.appendEvents(ctx, []*SessionEvent[*schema.Message]{event}) + require.NoError(t, err) + require.Len(t, store.appendEvents, 1) + assert.Equal(t, "sid", store.appendSessionIDs[0]) + + err = opened.handle.appendEvents(ctx, nil) + require.NoError(t, err) + require.Len(t, store.appendEvents, 2) + assert.Equal(t, "sid", store.appendSessionIDs[1]) + + require.NoError(t, opened.handle.close(ctx)) + require.NoError(t, opened.handle.close(ctx)) + err = opened.handle.appendEvents(ctx, nil) + require.ErrorIs(t, err, ErrSessionBusy) + + reopened, err := openLocalSession[*schema.Message](ctx, store, &openSessionRequest{sessionID: "sid"}) + require.NoError(t, err) + require.NoError(t, reopened.handle.close(ctx)) + + store.loadErr = errors.New("load failed") + opened, err = openLocalSession[*schema.Message](ctx, store, &openSessionRequest{sessionID: "sid-load-err"}) + require.NoError(t, err) + _, err = opened.handle.loadEvents(ctx, &LoadSessionEventsRequest{}) + require.ErrorContains(t, err, "load failed") + require.NoError(t, opened.handle.close(ctx)) + + store.loadErr = nil + store.appendErr = errors.New("append failed") + opened, err = openLocalSession[*schema.Message](ctx, store, &openSessionRequest{sessionID: "sid-append-err"}) + require.NoError(t, err) + err = opened.handle.appendEvents(ctx, []*SessionEvent[*schema.Message]{validTestPayload()}) + require.ErrorContains(t, err, "append failed") + require.NoError(t, opened.handle.close(ctx)) +} diff --git a/adk/failover_chatmodel.go b/adk/failover_chatmodel.go index f12d890bb..c09d41367 100644 --- a/adk/failover_chatmodel.go +++ b/adk/failover_chatmodel.go @@ -23,6 +23,8 @@ import ( "io" "log" + "github.com/google/uuid" + "github.com/cloudwego/eino/components" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/compose" @@ -56,24 +58,42 @@ func getFailoverHasMoreAttempts(ctx context.Context) bool { return v } +type failoverTimelineKey struct{} + +type failoverTimelineMeta struct { + ParentSpanID string + Attempt int +} + +func withFailoverTimeline(ctx context.Context, parentSpanID string, attempt int) context.Context { + return context.WithValue(ctx, failoverTimelineKey{}, failoverTimelineMeta{ParentSpanID: parentSpanID, Attempt: attempt}) +} + +func getFailoverTimeline(ctx context.Context) (failoverTimelineMeta, bool) { + v, ok := ctx.Value(failoverTimelineKey{}).(failoverTimelineMeta) + return v, ok +} + type typedFailoverProxyModel[M MessageType] struct { } -func (m *typedFailoverProxyModel[M]) prepareTarget(ctx context.Context) (model.BaseModel[M], error) { +func (m *typedFailoverProxyModel[M]) prepareTarget(ctx context.Context) (model.BaseModel[M], string, error) { target, ok := typedGetFailoverCurrentModel[M](ctx) if !ok { - return nil, errors.New("failover current model not found in context") + return nil, "", errors.New("failover current model not found in context") } + targetType, _ := components.GetType(target) + if !components.IsCallbacksEnabled(target) { target = typedCallbackInjectionModelWrapper[M]{}.wrapModel(target) } - return target, nil + return target, targetType, nil } func (m *typedFailoverProxyModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { - target, err := m.prepareTarget(ctx) + target, _, err := m.prepareTarget(ctx) if err != nil { var zero M return zero, err @@ -83,7 +103,7 @@ func (m *typedFailoverProxyModel[M]) Generate(ctx context.Context, input []M, op } func (m *typedFailoverProxyModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { - target, err := m.prepareTarget(ctx) + target, _, err := m.prepareTarget(ctx) if err != nil { return nil, err } @@ -230,6 +250,8 @@ func (f *failoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts var lastOutputMessage M var lastErr error + parentSpanID := uuid.NewString() + timelineAttempt := 1 // Try lastSuccessModel first if available. if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil { @@ -240,6 +262,7 @@ func (f *failoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess) modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0) + modelCtx = withFailoverTimeline(modelCtx, parentSpanID, timelineAttempt) result, err := f.inner.Generate(modelCtx, input, opts...) if err == nil { return result, nil @@ -252,7 +275,9 @@ func (f *failoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts return result, err } + emitFailoverRetryingTimeline[M](ctx, err) log.Printf("failover ChatModel.Generate lastSuccessModel failed: %v", err) + timelineAttempt++ } for attempt := uint(1); attempt <= f.config.MaxRetries; attempt++ { @@ -284,6 +309,7 @@ func (f *failoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts modelCtx := typedSetFailoverCurrentModel(ctx, currentModel) modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries) + modelCtx = withFailoverTimeline(modelCtx, parentSpanID, timelineAttempt) result, err := f.inner.Generate(modelCtx, currentInput, opts...) lastOutputMessage = result lastErr = err @@ -298,10 +324,13 @@ func (f *failoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts } if attempt < f.config.MaxRetries { + emitFailoverRetryingTimeline[M](ctx, err) log.Printf("failover ChatModel.Generate attempt %d failed: %v", attempt, err) } + timelineAttempt++ } + emitFailoverExhaustedTimeline[M](ctx, lastErr) return lastOutputMessage, lastErr } @@ -314,6 +343,8 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. var lastOutputMessage M var lastErr error + parentSpanID := uuid.NewString() + timelineAttempt := 1 // Try lastSuccessModel first if available. if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil { @@ -323,6 +354,7 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess) modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0) + modelCtx = withFailoverTimeline(modelCtx, parentSpanID, timelineAttempt) stream, err := f.inner.Stream(modelCtx, input, opts...) if err != nil { lastErr = err @@ -330,7 +362,9 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. if !f.needFailover(ctx, zero, err) { return nil, err } + emitFailoverRetryingTimeline[M](ctx, err) log.Printf("failover ChatModel.Stream lastSuccessModel failed: %v", err) + timelineAttempt++ } else { copies := stream.Copy(2) checkCopy := copies[0] @@ -345,7 +379,9 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. if !f.needFailover(ctx, outMsg, streamErr) { return nil, streamErr } + emitFailoverRetryingTimeline[M](ctx, streamErr) log.Printf("failover ChatModel.Stream lastSuccessModel failed: %v", streamErr) + timelineAttempt++ } else { return returnCopy, nil } @@ -378,6 +414,7 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. modelCtx := typedSetFailoverCurrentModel(ctx, currentModel) modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries) + modelCtx = withFailoverTimeline(modelCtx, parentSpanID, timelineAttempt) stream, err := f.inner.Stream(modelCtx, currentInput, opts...) if err != nil { lastErr = err @@ -389,8 +426,10 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. } if attempt < f.config.MaxRetries { + emitFailoverRetryingTimeline[M](ctx, err) log.Printf("failover ChatModel.Stream attempt %d failed: %v", attempt, err) } + timelineAttempt++ continue } @@ -425,8 +464,10 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. } if attempt < f.config.MaxRetries { + emitFailoverRetryingTimeline[M](ctx, streamErr) log.Printf("failover ChatModel.Stream attempt %d failed: %v", attempt, streamErr) } + timelineAttempt++ continue } @@ -434,9 +475,37 @@ func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts .. return returnCopy, nil } + emitFailoverExhaustedTimeline[M](ctx, lastErr) return nil, lastErr } +func emitFailoverRetryingTimeline[M MessageType](ctx context.Context, err error) { + sendSessionTimelineEvent(ctx, &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventSessionError, + Error: &SessionErrorEvent{ + Type: SessionErrorTypeModelFailover, + Message: timelineErrorMessage(err, nil), + RetryStatus: &RetryStatus{Type: "retrying"}, + }, + }) +} + +func emitFailoverExhaustedTimeline[M MessageType](ctx context.Context, err error) { + if err == nil { + return + } + sendSessionTimelineEvent(ctx, &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventSessionError, + Error: &SessionErrorEvent{ + Type: SessionErrorTypeModelFailover, + Message: timelineErrorMessage(err, nil), + RetryStatus: &RetryStatus{Type: "exhausted"}, + }, + }) +} + func typedConsumeStream[M MessageType](stream *schema.StreamReader[M]) (M, error) { var zero M defer stream.Close() diff --git a/adk/failover_chatmodel_test.go b/adk/failover_chatmodel_test.go index 8b8ca579b..6a39d40b3 100644 --- a/adk/failover_chatmodel_test.go +++ b/adk/failover_chatmodel_test.go @@ -48,6 +48,10 @@ func (m *fakeChatModel) IsCallbacksEnabled() bool { return m.callbacksEnabled } +func (m *fakeChatModel) GetType() string { + return "fake_chat_model" +} + func drainMessageStream(sr *schema.StreamReader[*schema.Message]) ([]*schema.Message, error) { defer sr.Close() var out []*schema.Message diff --git a/adk/filesystem/backend.go b/adk/filesystem/backend.go index 62ebee870..8f555cd44 100644 --- a/adk/filesystem/backend.go +++ b/adk/filesystem/backend.go @@ -158,6 +158,14 @@ type WriteRequest struct { Content string } +// AppendRequest contains parameters for appending content to a file. +type AppendRequest struct { + // FilePath is the path of the file to append to. + FilePath string + // Content is the data to append at the end of the file. + Content string +} + // EditRequest contains parameters for editing file content. type EditRequest struct { // FilePath is the path of the file to edit. @@ -236,6 +244,15 @@ type MultiModalReader interface { MultiModalRead(ctx context.Context, req *MultiModalReadRequest) (*MultiFileContent, error) } +// Appender appends content to the end of a file without rewriting the whole file, +// enabling efficient incremental writes — e.g. streaming a long-running background +// task's output to its output file as chunks arrive. +type Appender interface { + // Append adds req.Content to the end of the file at req.FilePath, creating the + // file if it does not exist. + Append(ctx context.Context, req *AppendRequest) error +} + // Backend is a pluggable, unified file backend protocol interface. // // All methods use struct-based parameters to allow future extensibility @@ -283,9 +300,12 @@ type Backend interface { } // ExecuteRequest contains parameters for executing a command. +// +// Foreground/background switching and timeouts are the caller's concern (e.g. the +// backgroundtask Manager): a backend simply runs the command and must honor ctx +// cancellation, which is how a timed-out or canceled run is stopped. type ExecuteRequest struct { - Command string // The command to execute - RunInBackendGround bool + Command string // The command to execute } // ExecuteResponse contains the response result of command execution. @@ -295,10 +315,16 @@ type ExecuteResponse struct { Truncated bool // Whether the output was truncated } +// Shell executes shell commands. Execute must honor ctx cancellation by stopping +// the underlying command (e.g. via exec.CommandContext): a timed-out or canceled +// run is stopped solely by canceling ctx, so an implementation that ignores it +// will leak the process and its goroutine after the run is reported stopped. type Shell interface { Execute(ctx context.Context, input *ExecuteRequest) (result *ExecuteResponse, err error) } +// StreamingShell is the streaming counterpart of Shell. ExecuteStreaming must honor +// ctx cancellation by stopping the underlying command, as described on Shell. type StreamingShell interface { ExecuteStreaming(ctx context.Context, input *ExecuteRequest) (result *schema.StreamReader[*ExecuteResponse], err error) } diff --git a/adk/filesystem/backend_inmemory.go b/adk/filesystem/backend_inmemory.go index 1a8118132..0403cb9d8 100644 --- a/adk/filesystem/backend_inmemory.go +++ b/adk/filesystem/backend_inmemory.go @@ -640,6 +640,26 @@ func (b *InMemoryBackend) Write(ctx context.Context, req *WriteRequest) error { return nil } +// Append adds content to the end of a file, creating it if it does not exist. +// It implements the optional Appender interface, letting OutputWriter stream +// task output incrementally without rewriting the whole file each time. +func (b *InMemoryBackend) Append(ctx context.Context, req *AppendRequest) error { + b.mu.Lock() + defer b.mu.Unlock() + + filePath := normalizePath(req.FilePath) + if entry, ok := b.files[filePath]; ok { + entry.content += req.Content + entry.modifiedAt = time.Now() + return nil + } + b.files[filePath] = &fileEntry{ + content: req.Content, + modifiedAt: time.Now(), + } + return nil +} + // Edit replaces string occurrences in a file. func (b *InMemoryBackend) Edit(ctx context.Context, req *EditRequest) error { b.mu.Lock() diff --git a/adk/handler.go b/adk/handler.go index f95244162..abe4ddac4 100644 --- a/adk/handler.go +++ b/adk/handler.go @@ -84,7 +84,7 @@ type ModelContext = TypedModelContext[*schema.Message] // Handlers can modify Instruction, Tools, and ReturnDirectly to customize agent behavior. // // This type is specific to ChatModelAgent. Other agent types may define their own context types. -type ChatModelAgentContext struct { +type ChatModelAgentContext[M MessageType] struct { // Instruction is the current instruction for the Agent execution. // It includes the instruction configured for the agent, additional instructions appended by framework // and AgentMiddleware, and modifications applied by previous BeforeAgent handlers. @@ -92,6 +92,8 @@ type ChatModelAgentContext struct { // to be (optionally) formatted with SessionValues and converted to system message. Instruction string + AgentInput *TypedAgentInput[M] + // Tools are the raw tools (without any wrapper or tool middleware) currently configured for the Agent execution. // They includes tools passed in AgentConfig, implicit tools added by framework such as transfer / exit tools, // and other tools already added by middlewares. @@ -139,7 +141,7 @@ type ChatModelAgentContext struct { type TypedChatModelAgentMiddleware[M MessageType] interface { // BeforeAgent is called before each agent run, allowing modification of // the agent's instruction and tools configuration. - BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) + BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[M]) (context.Context, *ChatModelAgentContext[M], error) // AfterAgent is called after the agent run reaches a successful terminal state. // Successful terminal states are: final answer (model response with no tool calls), @@ -296,7 +298,7 @@ func (b *TypedBaseChatModelAgentMiddleware[M]) WrapModel(_ context.Context, m mo return m, nil } -func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[M]) (context.Context, *ChatModelAgentContext[M], error) { return ctx, runCtx, nil } @@ -398,30 +400,36 @@ func DeleteRunLocalValue(ctx context.Context, key string) error { // TypedSendEvent sends a custom TypedAgentEvent to the event stream during agent execution. // This allows TypedChatModelAgentMiddleware implementations to emit custom events that will be // received by the caller iterating over the agent's event stream. +// To emit custom session timeline events during a Runner run, wrap a SessionEvent +// with Extension set and an x.* Kind in TypedAgentEvent.SessionEventVariant.Event. This is the +// canonical in-run path because Runner materializes identity, emits the live +// event, and persists it through the ordered session event pipeline. // -// Note: TypedSendEvent is a pure transport — it does NOT auto-assign message IDs. -// Framework-created messages (model output, tool results) receive IDs automatically -// via internal wrapper layers. If your middleware constructs its own messages, call -// EnsureMessageID before sending to assign an ID. +// TypedSendEvent assigns message IDs for message-bearing events before enqueueing +// them. Middleware authors only need to call EnsureMessageID directly when they +// need the ID before emitting the event, for example to build another event that +// references the message by ID. // -// This function can only be called from within a TypedChatModelAgentMiddleware during agent execution. -// Returns an error if called outside of an agent execution context. +// When called outside of an agent execution context, or from a path without an +// event generator, this function is a no-op. func TypedSendEvent[M MessageType](ctx context.Context, event *TypedAgentEvent[M]) error { execCtx := getTypedChatModelAgentExecCtx[M](ctx) if execCtx == nil || execCtx.generator == nil { - return fmt.Errorf("TypedSendEvent failed: must be called within a ChatModelAgent Run() or Resume() execution context") + return nil } - execCtx.send(event) + execCtx.send(ctx, event) return nil } // SendEvent sends a custom AgentEvent to the event stream during agent execution. // This allows ChatModelAgentMiddleware implementations to emit custom events that will be // received by the caller iterating over the agent's event stream. +// For custom session timeline events during a Runner run, set AgentEvent.SessionEventVariant.Event +// to an extension SessionEvent with an x.* Kind and send it through this function. // -// This function can only be called from within a ChatModelAgentMiddleware during agent execution. -// Returns an error if called outside of an agent execution context. +// When called outside of an agent execution context, or from a path without an +// event generator, this function is a no-op. func SendEvent(ctx context.Context, event *AgentEvent) error { return TypedSendEvent(ctx, event) } diff --git a/adk/handler_test.go b/adk/handler_test.go index 811cd2b27..cd304cbce 100644 --- a/adk/handler_test.go +++ b/adk/handler_test.go @@ -37,7 +37,7 @@ type testInstructionHandler struct { text string } -func (h *testInstructionHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (h *testInstructionHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { if runCtx.Instruction == "" { runCtx.Instruction = h.text } else if h.text != "" { @@ -51,7 +51,7 @@ type testInstructionFuncHandler struct { fn func(ctx context.Context, instruction string) (context.Context, string, error) } -func (h *testInstructionFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (h *testInstructionFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { newCtx, newInstruction, err := h.fn(ctx, runCtx.Instruction) if err != nil { return ctx, runCtx, err @@ -65,7 +65,7 @@ type testToolsHandler struct { tools []tool.BaseTool } -func (h *testToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (h *testToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { runCtx.Tools = append(runCtx.Tools, h.tools...) return ctx, runCtx, nil } @@ -75,7 +75,7 @@ type testToolsFuncHandler struct { fn func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error) } -func (h *testToolsFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (h *testToolsFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { newCtx, newTools, newReturnDirectly, err := h.fn(ctx, runCtx.Tools, runCtx.ReturnDirectly) if err != nil { return ctx, runCtx, err @@ -87,10 +87,10 @@ func (h *testToolsFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatMode type testBeforeAgentHandler struct { *BaseChatModelAgentMiddleware - fn func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) + fn func(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) } -func (h *testBeforeAgentHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (h *testBeforeAgentHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { return h.fn(ctx, runCtx) } @@ -894,10 +894,10 @@ func TestContextPropagation(t *testing.T) { Description: "Test agent", Model: cm, Handlers: []ChatModelAgentMiddleware{ - &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { return context.WithValue(ctx, key1, "value1"), runCtx, nil }}, - &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { handler2ReceivedValue = ctx.Value(key1) return ctx, runCtx, nil }}, @@ -962,7 +962,7 @@ func TestHandlerErrorHandling(t *testing.T) { Description: "Test agent", Model: cm, Handlers: []ChatModelAgentMiddleware{ - &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { return ctx, runCtx, assert.AnError }}, }, @@ -1042,7 +1042,7 @@ type countingHandler struct { mu sync.Mutex } -func (h *countingHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (h *countingHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { h.mu.Lock() h.beforeAgentCount++ h.mu.Unlock() diff --git a/adk/integration_middleware_test.go b/adk/integration_middleware_test.go new file mode 100644 index 000000000..7e37827c1 --- /dev/null +++ b/adk/integration_middleware_test.go @@ -0,0 +1,504 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk_test + +import ( + "context" + "fmt" + "os" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/middlewares/agentsmd" + "github.com/cloudwego/eino/adk/middlewares/dynamictool/toolsearch" + "github.com/cloudwego/eino/adk/middlewares/patchtoolcalls" + "github.com/cloudwego/eino/adk/middlewares/reduction" + "github.com/cloudwego/eino/adk/session" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// stubChatModel returns a fixed final assistant message and stops the React loop. +type stubChatModel struct { + reply string +} + +func (m *stubChatModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage(m.reply, nil), nil +} + +func (m *stubChatModel) Stream(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(m.reply, nil)}), nil +} + +// memBackend is a minimal Agents.md backend that serves an in-memory file. +type memBackend struct { + files map[string]string +} + +func (b *memBackend) Read(_ context.Context, req *filesystem.ReadRequest) (*filesystem.FileContent, error) { + content, ok := b.files[req.FilePath] + if !ok { + return nil, fmt.Errorf("file not found: %s: %w", req.FilePath, os.ErrNotExist) + } + return &filesystem.FileContent{Content: content}, nil +} + +// TestAgentsMDIntegration_PersistsMessageInserted is a true end-to-end test: +// it runs a real ChatModelAgent with the real agentsmd middleware through the +// Runner with session mode enabled, and verifies that the persistent event log +// contains a MessageInserted event carrying the agentsmd content. This covers +// the evaluation's "real middleware event emission" gap. +func TestAgentsMDIntegration_PersistsMessageInserted(t *testing.T) { + ctx := context.Background() + + backend := &memBackend{files: map[string]string{ + "AGENTS.md": "you are a careful agent", + }} + mw, err := agentsmd.New(ctx, &agentsmd.Config{ + Backend: backend, + AgentsMDFiles: []string{"AGENTS.md"}, + }) + require.NoError(t, err) + + model := &stubChatModel{reply: "ok"} + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "agentsmd-integration", + Description: "agentsmd integration test agent", + Instruction: "you are a test agent", + Model: model, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + require.NoError(t, err) + + store := session.NewInMemoryStore[*schema.Message](nil) + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "agentsmd-test", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hello") + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + // Read the persisted event log. + res, err := store.LoadEvents(ctx, "agentsmd-test", &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + + var sawInsertedAgentsmd bool + for _, se := range res.Events { + if se.MessageInserted == nil { + continue + } + ins := se.MessageInserted.Message + // The inserted message must carry the agentsmd marker so the next turn skips re-insertion. + if ins != nil && ins.Extra != nil { + if v, ok := ins.Extra["__agentsmd_content__"]; ok { + if b, ok := v.(bool); ok && b { + sawInsertedAgentsmd = true + assert.Contains(t, ins.Content, "you are a careful agent", + "persisted MessageInserted must carry the loaded agentsmd content") + } + } + } + } + assert.True(t, sawInsertedAgentsmd, + "agentsmd middleware running through ChatModelAgent + Runner must persist a MessageInserted event with the marker") +} + +// TestAgentsMDIntegration_NextTurnSkipsReinsertion verifies the prompt-cache +// stability invariant: after the first turn persists the agentsmd MessageInserted +// event, the second turn boots from the persisted state and the middleware does +// NOT re-insert (so the prefix bytes stay byte-identical). +func TestAgentsMDIntegration_NextTurnSkipsReinsertion(t *testing.T) { + ctx := context.Background() + + backend := &memBackend{files: map[string]string{ + "AGENTS.md": "stable agents.md prefix", + }} + mw, err := agentsmd.New(ctx, &agentsmd.Config{ + Backend: backend, + AgentsMDFiles: []string{"AGENTS.md"}, + }) + require.NoError(t, err) + + model := &stubChatModel{reply: "ok"} + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "agentsmd-stable", + Description: "agentsmd stable-prefix test", + Instruction: "test agent", + Model: model, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + require.NoError(t, err) + + store := session.NewInMemoryStore[*schema.Message](nil) + sid := "agentsmd-stable-session" + + // Turn 1. + runner1 := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + for it := runner1.Query(ctx, "first"); ; { + ev, ok := it.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + // Count agentsmd MessageInserted events after turn 1. + countAgentsmdInserts := func() int { + res, err := store.LoadEvents(ctx, sid, &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + count := 0 + for _, se := range res.Events { + if se.MessageInserted == nil { + continue + } + if se.MessageInserted.Message != nil && se.MessageInserted.Message.Extra != nil { + if v, ok := se.MessageInserted.Message.Extra["__agentsmd_content__"]; ok { + if b, ok := v.(bool); ok && b { + count++ + } + } + } + } + return count + } + require.Equal(t, 1, countAgentsmdInserts(), "first turn must insert exactly once") + + // Turn 2. + runner2 := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + for it := runner2.Query(ctx, "second"); ; { + ev, ok := it.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + // Critical assertion: still exactly one — turn 2 must NOT have inserted + // another agentsmd message because the marker is in the reconstructed history. + assert.Equal(t, 1, countAgentsmdInserts(), + "turn 2 must skip agentsmd re-insertion (prompt-cache prefix stability)") +} + +// dummyDynamicTool is a no-op dynamic tool the toolsearch middleware can advertise. +type dummyDynamicTool struct { + name string + desc string +} + +func (t *dummyDynamicTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: t.desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "q": {Type: schema.String, Desc: "q", Required: true}, + }), + }, nil +} + +func (t *dummyDynamicTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return `{"ok":true}`, nil +} + +// TestToolSearchIntegration_PersistsMessageInserted is the toolsearch +// counterpart of the agentsmd integration test: a real ChatModelAgent + real +// toolsearch middleware + Runner + InMemoryStore. It verifies the toolsearch +// reminder is persisted as a MessageInserted event and survives across turns. +func TestToolSearchIntegration_PersistsMessageInserted(t *testing.T) { + ctx := context.Background() + + mw, err := toolsearch.New(ctx, &toolsearch.Config{ + DynamicTools: []tool.BaseTool{ + &dummyDynamicTool{name: "weather", desc: "get weather for a city"}, + &dummyDynamicTool{name: "stocks", desc: "get a stock quote"}, + }, + }) + require.NoError(t, err) + + model := &stubChatModel{reply: "ok"} + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "toolsearch-integration", + Description: "toolsearch integration test agent", + Instruction: "test agent", + Model: model, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + require.NoError(t, err) + + store := session.NewInMemoryStore[*schema.Message](nil) + sid := "toolsearch-test" + sessionStore := store + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: sessionStore, + }) + + for it := runner.Query(ctx, "anything"); ; { + ev, ok := it.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + res, err := store.LoadEvents(ctx, sid, &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + + var sawInsertedReminder bool + for _, se := range res.Events { + if se.MessageInserted == nil { + continue + } + ins := se.MessageInserted.Message + if ins != nil && ins.Extra != nil { + if v, ok := ins.Extra["__toolsearch_reminder__"]; ok { + if b, ok := v.(bool); ok && b { + sawInsertedReminder = true + } + } + } + } + assert.True(t, sawInsertedReminder, + "toolsearch middleware running through ChatModelAgent + Runner must persist a MessageInserted event with the reminder marker") +} + +// TestPatchToolCallsIntegration_PersistsMessageInserted seeds the session event +// log with an assistant message that has a dangling tool call (no following +// tool result). On the next Run, the reconstructed history contains the +// dangling call; patchtoolcalls' BeforeModelRewriteState patches it by inserting +// a synthetic tool message and emitting a MessageInserted event. We verify the +// event reaches the persistent log. +func TestPatchToolCallsIntegration_PersistsMessageInserted(t *testing.T) { + ctx := context.Background() + + store := session.NewInMemoryStore[*schema.Message](nil) + sessionStore := store + sid := "patchtoolcalls-test" + + // Seed: an assistant message with a tool call but no corresponding tool result. + dangling := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call-1", + Type: "function", + Function: schema.FunctionCall{Name: "weather", Arguments: `{"city":"sf"}`}, + }, + }, + Extra: map[string]any{"_eino_msg_id": "dangling-msg-id"}, + } + user := &schema.Message{ + Role: schema.User, + Content: "what's the weather?", + Extra: map[string]any{"_eino_msg_id": "user-msg-id"}, + } + + for _, m := range []*schema.Message{user, dangling} { + se := &adk.SessionEvent[*schema.Message]{EventID: uuid.NewString(), Kind: adk.SessionEventMessage, Message: m} + err := store.AppendEvents(ctx, sid, []*adk.SessionEvent[*schema.Message]{se}) + require.NoError(t, err) + } + + // Wire patchtoolcalls into a ChatModelAgent. + mw, err := patchtoolcalls.New(ctx, nil) + require.NoError(t, err) + + model := &stubChatModel{reply: "done"} + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "patchtoolcalls-integration", + Description: "patchtoolcalls integration test agent", + Instruction: "test agent", + Model: model, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + require.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: sessionStore, + }) + + for it := runner.Query(ctx, "go"); ; { + ev, ok := it.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + // Read events back; among the events appended on this turn there should be + // a MessageInserted carrying a Tool-role synthetic message. + res, err := store.LoadEvents(ctx, sid, &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + var sawInsertedToolResult bool + for _, se := range res.Events { + if se.MessageInserted == nil { + continue + } + ins := se.MessageInserted.Message + if ins != nil && ins.Role == schema.Tool && ins.ToolCallID == "call-1" { + sawInsertedToolResult = true + } + } + assert.True(t, sawInsertedToolResult, + "patchtoolcalls middleware must persist a MessageInserted event for the synthetic tool result") +} + +// TestReductionIntegration_PersistsBothMessageUpdated seeds the session log +// with two rounds of (assistant tool call → tool result). With reduction's +// ClearRetentionSuffixLimit=1 (the framework default), the LAST round is +// retained and the FIRST round is cleared. Reduction emits MessageUpdated for +// both the assistant tool-call message (args replaced + cleared flag) and the +// tool-result message (content replaced). Both must reach the persistent log. +func TestReductionIntegration_PersistsBothMessageUpdated(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + sessionStore := store + sid := "reduction-test" + + // Seed the session: user → assistant call A → tool result A → assistant call B → tool result B. + // With ClearRetentionSuffixLimit=1, round B is retained; round A is cleared. + user := &schema.Message{ + Role: schema.User, + Content: "do the thing", + Extra: map[string]any{"_eino_msg_id": "user-id"}, + } + assistantA := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "tc-A", Type: "function", Function: schema.FunctionCall{Name: "noop", Arguments: `{"q":"A"}`}}, + }, + Extra: map[string]any{"_eino_msg_id": "assistant-A-id"}, + } + toolResultA := &schema.Message{ + Role: schema.Tool, + ToolCallID: "tc-A", + ToolName: "noop", + Content: "raw content A", + Extra: map[string]any{"_eino_msg_id": "tool-A-id"}, + } + assistantB := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "tc-B", Type: "function", Function: schema.FunctionCall{Name: "noop", Arguments: `{"q":"B"}`}}, + }, + Extra: map[string]any{"_eino_msg_id": "assistant-B-id"}, + } + toolResultB := &schema.Message{ + Role: schema.Tool, + ToolCallID: "tc-B", + ToolName: "noop", + Content: "raw content B", + Extra: map[string]any{"_eino_msg_id": "tool-B-id"}, + } + for _, m := range []*schema.Message{user, assistantA, toolResultA, assistantB, toolResultB} { + se := &adk.SessionEvent[*schema.Message]{EventID: uuid.NewString(), Kind: adk.SessionEventMessage, Message: m} + err := store.AppendEvents(ctx, sid, []*adk.SessionEvent[*schema.Message]{se}) + require.NoError(t, err) + } + + // Reduction config: token counter always exceeds threshold; clear handler always clears. + mw, err := reduction.New(ctx, &reduction.Config{ + SkipTruncation: true, + TokenCounter: func(_ context.Context, _ []*schema.Message, _ []*schema.ToolInfo) (int64, error) { + return 1000000, nil + }, + MaxTokensForClear: 1, + // ClearRetentionSuffixLimit defaults to 1: round B is retained, round A is cleared. + GenClearOffloadFilePath: func(_ context.Context, td *reduction.ToolDetail) (string, error) { + return "/tmp/" + td.ToolContext.CallID, nil + }, + ToolConfig: map[string]*reduction.ToolReductionConfig{ + "noop": { + SkipClear: false, + ClearHandler: func(_ context.Context, _ *reduction.ToolDetail) (*reduction.ClearResult, error) { + return &reduction.ClearResult{ + NeedClear: true, + ToolArgument: &schema.ToolArgument{Text: `{"q":"[cleared]"}`}, + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "[cleared]"}}}, + }, nil + }, + }, + }, + }) + require.NoError(t, err) + + model := &stubChatModel{reply: "ok"} + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "reduction-integration", + Description: "reduction integration test agent", + Instruction: "test agent", + Model: model, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + require.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: sessionStore, + }) + + for it := runner.Query(ctx, "go"); ; { + ev, ok := it.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + res, err := store.LoadEvents(ctx, sid, &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + + var sawAssistantUpdated, sawToolUpdated bool + for _, se := range res.Events { + if se.MessageUpdated == nil { + continue + } + switch se.MessageUpdated.MessageID { + case "assistant-A-id": + sawAssistantUpdated = true + case "tool-A-id": + sawToolUpdated = true + } + } + assert.True(t, sawAssistantUpdated, + "reduction must emit MessageUpdated for the cleared assistant tool-call message (round A)") + assert.True(t, sawToolUpdated, + "reduction must emit MessageUpdated for the cleared tool-result message (round A)") +} diff --git a/adk/interface.go b/adk/interface.go index 8905950d9..e1ed5c481 100644 --- a/adk/interface.go +++ b/adk/interface.go @@ -20,14 +20,20 @@ import ( "bytes" "context" "encoding/gob" + "errors" "fmt" "io" + "time" "github.com/cloudwego/eino/components" "github.com/cloudwego/eino/internal/core" "github.com/cloudwego/eino/schema" ) +func newEventTimestamp() time.Time { + return time.Now().UTC() +} + // ComponentOfAgent is the component type identifier for ADK agents in callbacks. // Use this to filter callback events to only agent-related events. const ComponentOfAgent components.Component = "Agent" @@ -432,6 +438,12 @@ type TypedAgentEvent[M MessageType] struct { Action *AgentAction Err error + + // SessionEventVariant is the first-class live session envelope. Event carries + // a materialized durable SessionEvent. MessageStreamRef carries only the + // reserved durable metadata for a streaming message whose content remains in + // Output.MessageOutput.MessageStream. + SessionEventVariant *SessionEventVariant[M] } // AgentEvent is the default event type using *schema.Message. @@ -517,3 +529,55 @@ func concatMessageStream[M MessageType](stream *schema.StreamReader[M]) (M, erro panic("unreachable: unknown MessageType") } } + +func materializeMessageStreamPrefix[M MessageType](stream *schema.StreamReader[M]) (msg M, hasChunks bool, streamErr error, err error) { + var zero M + switch s := any(stream).(type) { + case *schema.StreamReader[*schema.Message]: + defer s.Close() + var msgs []*schema.Message + for { + frame, recvErr := s.Recv() + if errors.Is(recvErr, io.EOF) { + break + } + if recvErr != nil { + streamErr = recvErr + break + } + msgs = append(msgs, frame) + } + if len(msgs) == 0 { + return zero, false, streamErr, nil + } + result, concatErr := schema.ConcatMessages(msgs) + if concatErr != nil { + return zero, true, streamErr, concatErr + } + return any(result).(M), true, streamErr, nil + case *schema.StreamReader[*schema.AgenticMessage]: + defer s.Close() + var msgs []*schema.AgenticMessage + for { + frame, recvErr := s.Recv() + if errors.Is(recvErr, io.EOF) { + break + } + if recvErr != nil { + streamErr = recvErr + break + } + msgs = append(msgs, frame) + } + if len(msgs) == 0 { + return zero, false, streamErr, nil + } + result, concatErr := schema.ConcatAgenticMessages(msgs) + if concatErr != nil { + return zero, true, streamErr, concatErr + } + return any(result).(M), true, streamErr, nil + default: + panic("unreachable: unknown MessageType") + } +} diff --git a/adk/interrupt.go b/adk/interrupt.go index 3d31054f5..9d51108bc 100644 --- a/adk/interrupt.go +++ b/adk/interrupt.go @@ -49,6 +49,11 @@ type ResumeInfo struct { type InterruptInfo struct { Data any + // CheckPointID is the checkpoint key used to persist this interrupted run, + // when checkpoint persistence is enabled. Pass this ID to Runner.Resume or + // Runner.ResumeWithParams to continue the same suspended execution. + CheckPointID string + // InterruptContexts provides a structured, user-facing view of the interrupt chain. // Each context represents a step in the agent hierarchy that was interrupted. InterruptContexts []*InterruptCtx @@ -288,12 +293,32 @@ func runnerSaveCheckPointImpl( info *InterruptInfo, is *core.InterruptSignal, ) error { - if store == nil { + if isNilCheckPointStore(store) { return nil } - runCtx := getRunCtx(ctx) + data, err := encodeRunnerCheckPointImpl(enableStreaming, ctx, info, is) + if err != nil { + return err + } + return store.Set(ctx, key, data) +} + +func encodeRunnerCheckPointImpl( + enableStreaming bool, + ctx context.Context, + info *InterruptInfo, + is *core.InterruptSignal, +) ([]byte, error) { + return encodeRunnerCheckPointWithRunCtx(enableStreaming, getRunCtx(ctx), info, is) +} +func encodeRunnerCheckPointWithRunCtx( + enableStreaming bool, + runCtx *runContext, + info *InterruptInfo, + is *core.InterruptSignal, +) ([]byte, error) { id2Addr, id2State := core.SignalToPersistenceMaps(is) buf := &bytes.Buffer{} @@ -305,9 +330,9 @@ func runnerSaveCheckPointImpl( EnableStreaming: enableStreaming, }) if err != nil { - return fmt.Errorf("failed to encode checkpoint: %w", err) + return nil, fmt.Errorf("failed to encode checkpoint: %w", err) } - return store.Set(ctx, key, buf.Bytes()) + return buf.Bytes(), nil } const bridgeCheckpointID = "adk_react_mock_key" @@ -317,21 +342,26 @@ func newBridgeStore() *bridgeStore { } func newResumeBridgeStore(checkPointID string, data []byte) *bridgeStore { + payload := append([]byte{}, data...) return &bridgeStore{ - data: map[string][]byte{checkPointID: data}, + data: map[string][]byte{checkPointID: payload}, + lastKey: checkPointID, + lastPayload: payload, } } type bridgeStore struct { - mu sync.Mutex - data map[string][]byte + mu sync.Mutex + data map[string][]byte + lastKey string + lastPayload []byte } func (m *bridgeStore) Get(_ context.Context, key string) ([]byte, bool, error) { m.mu.Lock() defer m.mu.Unlock() if v, ok := m.data[key]; ok { - return v, true, nil + return append([]byte{}, v...), true, nil } return nil, false, nil } @@ -342,10 +372,22 @@ func (m *bridgeStore) Set(_ context.Context, key string, checkPoint []byte) erro if m.data == nil { m.data = make(map[string][]byte) } - m.data[key] = checkPoint + payload := append([]byte{}, checkPoint...) + m.data[key] = payload + m.lastKey = key + m.lastPayload = payload return nil } +func (m *bridgeStore) LastCheckpoint() (key string, payload []byte, ok bool) { + m.mu.Lock() + defer m.mu.Unlock() + if m.lastKey == "" { + return "", nil, false + } + return m.lastKey, append([]byte{}, m.lastPayload...), true +} + func getNextResumeAgent(ctx context.Context, _ *ResumeInfo) (string, error) { nextAgents, err := core.GetNextResumptionPoints(ctx) if err != nil { diff --git a/adk/interrupt_test.go b/adk/interrupt_test.go index 480c0f8f7..ec57df580 100644 --- a/adk/interrupt_test.go +++ b/adk/interrupt_test.go @@ -59,7 +59,7 @@ func TestPreprocessADKCheckpoint(t *testing.T) { }) } -func (h *interruptTestToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { +func (h *interruptTestToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext[*schema.Message]) (context.Context, *ChatModelAgentContext[*schema.Message], error) { runCtx.Tools = append(runCtx.Tools, h.tools...) return ctx, runCtx, nil } @@ -991,8 +991,22 @@ func TestWorkflowInterrupt(t *testing.T) { }, } - assert.Contains(t, events, parallelMessageEvents[0]) - assert.Contains(t, events, parallelMessageEvents[1]) + assertParallelMessageEvent := func(want *AgentEvent) { + t.Helper() + for _, event := range events { + if event.AgentName == want.AgentName && + assert.ObjectsAreEqual(want.RunPath, event.RunPath) && + event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.Message != nil && + event.Output.MessageOutput.Message.Content == want.Output.MessageOutput.Message.Content { + return + } + } + assert.Failf(t, "missing parallel message event", "want=%v", want) + } + + assertParallelMessageEvent(parallelMessageEvents[0]) + assertParallelMessageEvent(parallelMessageEvents[1]) assert.NotNil(t, interruptEvent) assert.Equal(t, "parallel agent", interruptEvent.AgentName) diff --git a/adk/message_id_test.go b/adk/message_id_test.go index ef8533575..fe2c300cc 100644 --- a/adk/message_id_test.go +++ b/adk/message_id_test.go @@ -470,11 +470,10 @@ func TestMessageID_UserInputNoAutoID(t *testing.T) { } } -// Scenario 8: Middleware must call EnsureMessageID before SendEvent; pointer identity ensures state consistency -// TestMessageID_SendEvent_MiddlewareMustEnsureID verifies that TypedSendEvent is a pure -// transport and does NOT auto-assign message IDs. Middleware authors must call -// EnsureMessageID themselves before sending. -func TestMessageID_SendEvent_MiddlewareMustEnsureID(t *testing.T) { +// Scenario 8: SendEvent assigns message IDs before enqueue; pointer identity ensures state consistency. +// TestMessageID_SendEvent_AutoEnsuresID verifies that middleware-created messages +// receive IDs at the SendEvent boundary. +func TestMessageID_SendEvent_AutoEnsuresID(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -498,21 +497,18 @@ func TestMessageID_SendEvent_MiddlewareMustEnsureID(t *testing.T) { // Middleware creates a new message and writes the SAME pointer to both state and event middlewareMsg = schema.AssistantMessage("middleware injected", nil) - // Middleware is responsible for assigning the ID before sending - EnsureMessageID(middlewareMsg) - // Write to state state.Messages = append(state.Messages, middlewareMsg) - // Send as event — TypedSendEvent does NOT auto-assign ID + // Send as event — TypedSendEvent assigns ID on the shared pointer. event := EventFromMessage(middlewareMsg, nil, schema.Assistant, "") err := SendEvent(ctx, event) if err != nil { return err } - // Because we called EnsureMessageID on the shared pointer, - // the state copy also has the ID (pointer identity) + // Because SendEvent ensures ID on the shared pointer, the state + // copy also has the ID (pointer identity). stateMsgIDAfterSendEvent = internal.GetMessageID(middlewareMsg.Extra) return nil @@ -538,10 +534,10 @@ func TestMessageID_SendEvent_MiddlewareMustEnsureID(t *testing.T) { // We expect at least 2 events: model response + middleware injected message require.GreaterOrEqual(t, len(allEvents), 2) - // The middleware message pointer should have an ID (assigned by middleware via EnsureMessageID) + // The middleware message pointer should have an ID assigned at SendEvent time. require.NotNil(t, middlewareMsg) middlewareMsgID := GetMessageID(middlewareMsg) - assert.NotEmpty(t, middlewareMsgID, "middleware should have assigned an ID via EnsureMessageID") + assert.NotEmpty(t, middlewareMsgID, "SendEvent should assign an ID") assert.True(t, isValidUUID(middlewareMsgID)) // The ID captured right after SendEvent (via pointer identity) should be the same @@ -566,7 +562,7 @@ func TestMessageID_SendEvent_MiddlewareMustEnsureID(t *testing.T) { middlewareEventMsgID, middlewareMsgID) } -func TestAttack_ConcatCorruptsIDIfMultipleChunksCarryIt(t *testing.T) { +func TestMessageID_ConcatCorruptsIDIfMultipleChunksCarryIt(t *testing.T) { id := "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" msgs := []*schema.Message{ {Role: schema.Assistant, Content: "chunk1", Extra: map[string]any{internal.EinoMsgIDKey: id}}, @@ -583,7 +579,7 @@ func TestAttack_ConcatCorruptsIDIfMultipleChunksCarryIt(t *testing.T) { assert.Equal(t, "chunk1chunk2chunk3", concatenated.Content) } -func TestAttack_ConcatPreservesIDIfOnlyFirstChunkHasIt(t *testing.T) { +func TestMessageID_ConcatPreservesIDIfOnlyFirstChunkHasIt(t *testing.T) { id := "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" msgs := []*schema.Message{ {Role: schema.Assistant, Content: "chunk1", Extra: map[string]any{internal.EinoMsgIDKey: id}}, @@ -598,7 +594,7 @@ func TestAttack_ConcatPreservesIDIfOnlyFirstChunkHasIt(t *testing.T) { assert.Equal(t, "chunk1chunk2chunk3", concatenated.Content) } -func TestAttack_ConcurrentGenerate_NoSharedExtraMutation(t *testing.T) { +func TestMessageID_ConcurrentGenerateNoSharedExtraMutation(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -651,7 +647,7 @@ func TestAttack_ConcurrentGenerate_NoSharedExtraMutation(t *testing.T) { // The important thing is no panic and unique IDs } -func TestAttack_GenerateCopyDoesNotAffectOriginal(t *testing.T) { +func TestMessageID_GenerateCopyDoesNotAffectOriginal(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -870,7 +866,7 @@ func TestMessageID_AgenticPublicAPIHelpers(t *testing.T) { // TestAttack_PopToolMsgID_DoublePop tests that calling popToolMsgID twice for the // same key returns "" on second call. -func TestAttack_PopToolMsgID_DoublePop(t *testing.T) { +func TestMessageID_PopToolMsgIDDoublePop(t *testing.T) { st := &typedState[*schema.Message]{} st.setToolMsgID("myTool", "call-1", "uuid-abc") @@ -913,7 +909,7 @@ func (t *namedFakeToolForTest) InvokableRun(_ context.Context, _ string, _ ...to // TestAttack_ToolMsgIDConsistency_MultipleTools is an integration test: when an agent // has multiple tools called in one turn, verify that EACH tool's event message ID // matches its corresponding state message ID. -func TestAttack_ToolMsgIDConsistency_MultipleTools(t *testing.T) { +func TestMessageID_ToolMsgIDConsistencyMultipleTools(t *testing.T) { ctx := context.Background() ctrl := gomock.NewController(t) defer ctrl.Finish() @@ -1008,7 +1004,7 @@ func TestAttack_ToolMsgIDConsistency_MultipleTools(t *testing.T) { // TestAttack_ToolResultToBlocks_EdgeCases verifies toolResultToBlocks handles // nil ToolResult, empty Parts, and Parts with nil media fields. -func TestAttack_ToolResultToBlocks_EdgeCases(t *testing.T) { +func TestMessageID_ToolResultToBlocksEdgeCases(t *testing.T) { t.Run("nil ToolResult", func(t *testing.T) { blocks := toolResultToBlocks(nil) assert.Nil(t, blocks, "nil ToolResult should produce nil blocks") diff --git a/adk/middlewares/agentsmd/agentsmd.go b/adk/middlewares/agentsmd/agentsmd.go index 5339998b2..1ced2974c 100644 --- a/adk/middlewares/agentsmd/agentsmd.go +++ b/adk/middlewares/agentsmd/agentsmd.go @@ -15,9 +15,10 @@ */ // Package agentsmd provides a middleware that automatically injects Agents.md -// file contents into model input messages. The injection is transient — content -// is prepended at model call time and never persisted to conversation state, -// so it is naturally excluded from summarization / compression. +// file contents into model input messages. The injected message is appended to +// state.Messages once per session (idempotent via an Extra marker) and persisted +// in the session event log so subsequent turns reuse it without regenerating +// (which would invalidate the prompt cache). package agentsmd import ( @@ -111,10 +112,37 @@ func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state } nState := *state - nState.Messages = typedInsertBeforeFirstUser(state.Messages, content) + newMessages, insertedMsg, anchorMsg := typedInsertBeforeFirstUser(state.Messages, content) + nState.Messages = newMessages + + // Emit MessageInserted so the persisted event log reflects the inserted + // agentsmd message. On the next turn, reconstruction includes it, the + // idempotent marker suppresses re-insertion, and the prompt-cache prefix + // remains byte-identical. + var beforeID string + if !isNilMessage(anchorMsg) { + beforeID = adk.GetMessageID(anchorMsg) + } + _ = adk.TypedSendEvent(ctx, &adk.TypedAgentEvent[M]{ + SessionEventVariant: &adk.SessionEventVariant[M]{ + Event: &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessageInserted, + MessageInserted: &adk.MessageInsertedEvent[M]{ + Message: insertedMsg, + BeforeMessageID: beforeID, + }, + }, + }, + }) + return ctx, &nState, nil } +func isNilMessage[M adk.MessageType](msg M) bool { + var zero M + return any(msg) == any(zero) +} + // hasAgentsMDExtra checks whether a message has the agentsmd extra key set. func hasAgentsMDExtra[M adk.MessageType](msg M) bool { switch v := any(msg).(type) { @@ -134,20 +162,26 @@ func hasAgentsMDExtra[M adk.MessageType](msg M) bool { return false } -// typedInsertBeforeFirstUser inserts a user message with agentsmd content before the first User message. -func typedInsertBeforeFirstUser[M adk.MessageType](msgs []M, content string) []M { - newMsg := makeUserMsgWithExtra[M](content) - result := make([]M, 0, len(msgs)+1) +// typedInsertBeforeFirstUser inserts a user message with agentsmd content before +// the first User message. Returns the updated slice, the inserted message (with +// an assigned eino message ID), and the anchor message (the first user message +// the inserted message was placed before; zero value if no user message exists +// and the inserted message was appended at the end). +func typedInsertBeforeFirstUser[M adk.MessageType](msgs []M, content string) (result []M, insertedMsg M, anchorMsg M) { + insertedMsg = makeUserMsgWithExtra[M](content) + adk.EnsureMessageID(insertedMsg) + result = make([]M, 0, len(msgs)+1) for i, msg := range msgs { if isUserRole(msg) { - result = append(result, newMsg) + result = append(result, insertedMsg) result = append(result, msgs[i:]...) - return result + anchorMsg = msg + return result, insertedMsg, anchorMsg } result = append(result, msg) } - result = append(result, newMsg) - return result + result = append(result, insertedMsg) + return result, insertedMsg, anchorMsg } func isUserRole[M adk.MessageType](msg M) bool { diff --git a/adk/middlewares/automemory/automemory.go b/adk/middlewares/automemory/automemory.go new file mode 100644 index 000000000..3383e2766 --- /dev/null +++ b/adk/middlewares/automemory/automemory.go @@ -0,0 +1,1013 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package automemory provides middleware that injects and persists session +// memories around chat-model agent runs. +package automemory + +import ( + "context" + "fmt" + "path/filepath" + "sort" + "strings" + "sync" + + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/adk" + ainternal "github.com/cloudwego/eino/adk/middlewares/automemory/internal" + fsmw "github.com/cloudwego/eino/adk/middlewares/filesystem" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[*memoryExtra]("_eino_adk_automemory_extra") +} + +type Config[M adk.MessageType] struct { + // MemoryDirectory is the persistent memory root directory exposed to automemory. + // Required. Relative paths are resolved against the process working directory. + MemoryDirectory string + + // MemoryBackend is the storage backend used by MemoryDirectory. + // Required. File operations are bounded to MemoryDirectory. + MemoryBackend Backend + + // GenInstruction returns the runtime memory instruction appended to the main agent system prompt. + // Use it to customize how strongly the main agent should read from and write to memory during normal task execution. + // It does not control the post-run extraction agent; use Write.GenInstruction for extraction-specific save criteria. + // The framework always appends the memory directory manifest after this block. + // Optional. Defaults to the built-in auto memory instruction. + GenInstruction func(ctx context.Context) (string, error) + + // Model is the default model used by topic selection and memory extraction. + // Per-read/per-write overrides can be configured in Read.Model / Write.Model. + // Optional. Defaults to nil; topic selection and extraction must then provide their own models. + Model model.BaseModel[M] + + // Read controls how memories are loaded and injected. + // Optional. Defaults to Sync load with topic selection enabled (if Model is set). + Read *ReadConfig[M] + + // Write controls post-run memory extraction and persistence. + // Optional. Default: disabled. + Write *WriteConfig[M] + + // Coordination controls session identity and distributed async extraction coordination. + // Optional. Defaults to a local in-process coordinator. + Coordination *CoordinationConfig[M] + + // OnError is called when automemory encounters an error. Errors are best-effort by default: + // the middleware will skip memory injection and allow the agent to continue. + // Optional. + OnError func(ctx context.Context, stage ErrorStage, err error) +} + +type ReadMode string + +const ( + ReadModeSync ReadMode = "sync" + ReadModeAsync ReadMode = "async" +) + +type ReadConfig[M adk.MessageType] struct { + Mode ReadMode + + // Model is used for topic selection. Defaults to Config.Model. + Model model.BaseModel[M] + + // Index controls whether and how MEMORY.md is loaded as a memory index reminder. + // Optional. Defaults to enabled with MEMORY.md as the index file. + Index *IndexConfig + + // TopicSelection controls the "LLM select topics" path. + // Optional. If nil, default topic selection settings are applied. + // Topic selection becomes active when Read.Model is available. + TopicSelection *TopicSelectionConfig +} + +type IndexConfig struct { + // FileName is the index file name under MemoryDirectory. + // Optional. Defaults to MEMORY.md. + FileName string + + // MaxLines caps index content injected into the memory index reminder. + // Optional. Defaults to package default. + MaxLines int + + // MaxBytes caps index content injected into the memory index reminder. + // Optional. Defaults to package default. + MaxBytes int +} + +type TopicSelectionConfig struct { + // Enable controls whether topic memory selection is enabled. + // When false, automemory will not query, rank, read, or inject topic memories. + // Optional. Defaults to true when nil. + Enable *bool + + // CandidateGlob is matched against the RELATIVE path under MemoryDirectory. + // Example: "**/*.md" + // Optional. Defaults to CandidateGlobPattern. + CandidateGlob string + + // CandidateLimit caps the number of candidate topic files considered for selection. + // Optional. Defaults to 200. + CandidateLimit int + + // CandidatePreviewLines are read from each candidate to parse YAML frontmatter. + // Optional. Defaults to 30. + CandidatePreviewLines int + + // TopK caps the number of topic memory files selected for injection. + // Optional. Defaults to 5. + TopK int + + // MaxLines caps single topic memory file read lines. + // Optional. Defaults to 200. + MaxLines int + + // MaxBytes caps single topic memory file read bytes. + // Optional. Defaults to 4k. + MaxBytes int + + // MaxTotalBytes caps the total rendered topic memory reminder. + // Optional. Defaults to 16k. + MaxTotalBytes int +} + +type WriteMode string + +const ( + WriteModeDisabled WriteMode = "disabled" + WriteModeAsync WriteMode = "async" + WriteModeSync WriteMode = "sync" +) + +type WriteConfig[M adk.MessageType] struct { + Mode WriteMode + + // Model is used for memory extraction. Defaults to Config.Model. + Model model.BaseModel[M] + + // MaxTurns caps the extractor's tool-call loop. + MaxTurns int + + // GenInstruction returns the save policy block used by the post-run memory extraction agent. + // Use it to customize which observations should or should not be persisted after a run. + // This replaces the extractor prompt's built-in "What to save" and "What NOT to save" sections; runtime memory behavior + // in the main agent system prompt is controlled by Config.GenInstruction. + // Optional. Defaults to the built-in extraction save criteria. + GenInstruction func(ctx context.Context) (string, error) + + // HandleExtractionIterator, if set, is called with the extractionAgent's event + // iterator returned by Run(). The handler is responsible for draining the + // iterator (calling Next until it returns ok=false) and returning any error + // it wants to surface to the middleware. + // + // If nil, automemory uses the default drain behavior: ignore all events and + // return the first ev.Err encountered (if any). + HandleExtractionIterator func(ctx context.Context, iter *adk.AsyncIterator[*adk.TypedAgentEvent[M]]) error +} + +type middleware[M adk.MessageType] struct { + adk.TypedBaseChatModelAgentMiddleware[M] + + cfg *Config[M] + + resolvedMemoryDirectory string + boundedMemoryBackend *ainternal.FSBackend + + topicSelectionModel model.BaseModel[M] + extractionHandler adk.TypedChatModelAgentMiddleware[M] + topicSelectionTool *schema.ToolInfo + coordination *CoordinationConfig[M] +} + +type selectionFuture struct { + done chan struct{} + mu sync.Mutex + + // Store an immutable snapshot to avoid being mutated via shared pointers. + content string + err error + applied bool +} + +type ctxKeySelectionFuture struct{} + +const ( + memoryExtraKey = "__eino_automemory__" +) + +type memoryExtra struct { + Type string + Cursor int +} + +// New creates an automemory middleware from the provided configuration. +func New[M adk.MessageType](ctx context.Context, config *Config[M]) (adk.TypedChatModelAgentMiddleware[M], error) { + if config == nil { + return nil, fmt.Errorf("auto memory config: invalid") + } + + cfg := cloneConfig(config) + if cfg.MemoryDirectory == "" || cfg.MemoryBackend == nil { + return nil, fmt.Errorf("auto memory config: invalid") + } + + resolvedMemoryDir, err := ainternal.ResolveMemoryDir(cfg.MemoryDirectory) + if err != nil { + return nil, fmt.Errorf("auto memory config: resolve memory directory: %w", err) + } + boundedMemoryBackend, err := ainternal.NewFSBackend(cfg.MemoryBackend, ainternal.FSBackendConfig{ + BaseDir: resolvedMemoryDir, + NotFoundAsContent: true, + ErrorPrefix: "memory backend", + }) + if err != nil { + return nil, err + } + if cfg.Read == nil { + cfg.Read = &ReadConfig[M]{} + } + applyReadDefaults(cfg) + + m := &middleware[M]{ + TypedBaseChatModelAgentMiddleware: adk.TypedBaseChatModelAgentMiddleware[M]{}, + cfg: cfg, + resolvedMemoryDirectory: resolvedMemoryDir, + boundedMemoryBackend: boundedMemoryBackend, + coordination: cfg.Coordination, + } + + m.topicSelectionTool = topicSelectionToolInfo() + if topicSelectionConfigEnabled(cfg.Read.TopicSelection) && cfg.Read.Model != nil { + m.topicSelectionModel = &modelWithTools[M]{ + base: cfg.Read.Model, + tools: []*schema.ToolInfo{m.topicSelectionTool}, + } + } + + if cfg.Write.Mode != WriteModeDisabled && cfg.Write.Model != nil { + fileSystemMiddleware, err := fsmw.NewTyped[M](ctx, &fsmw.MiddlewareConfig{ + Backend: boundedMemoryBackend, + LsToolConfig: &fsmw.ToolConfig{Disable: true}, + GrepToolConfig: &fsmw.ToolConfig{Disable: true}, + }) + if err != nil { + return nil, err + } + m.extractionHandler = fileSystemMiddleware + } + + return m, nil +} + +func (m *middleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { + if runCtx == nil { + return ctx, runCtx, nil + } + nRunCtx := *runCtx + + // Sync distributed write cursor back into message extras so later runs on other + // machines still carry a transcript-local marker. + if nRunCtx.AgentInput != nil && len(nRunCtx.AgentInput.Messages) > 0 && m.coordination != nil && m.coordination.Coordinator != nil { + if sessionID, err := m.resolveSessionID(ctx, &adk.TypedChatModelAgentState[M]{Messages: nRunCtx.AgentInput.Messages}); err == nil && sessionID != "" { + localCursor := getWriteCursorFromMessages(nRunCtx.AgentInput.Messages) + coordKey := m.coordinatorKey(sessionID) + if remoteCursor, ok, err := getCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey); err == nil && ok && remoteCursor > localCursor { + st := markWriteCursor(&adk.TypedChatModelAgentState[M]{Messages: nRunCtx.AgentInput.Messages}, remoteCursor) + if st != nil { + nRunCtx.AgentInput = &adk.TypedAgentInput[M]{ + Messages: st.Messages, + EnableStreaming: nRunCtx.AgentInput.EnableStreaming, + } + } + } + } + } + + // 1) System prompt: inject stable auto memory instruction and directory manifest (best-effort). + instruction, err := m.renderInstruction(ctx, nRunCtx.Instruction) + if err != nil { + m.onErr(ctx, OnErrorStageRenderInstruction, err) + } else { + nRunCtx.Instruction = instruction + } + + if nRunCtx.AgentInput == nil || len(nRunCtx.AgentInput.Messages) == 0 { + return ctx, &nRunCtx, nil + } + + var reminders []M + + // 2) Memory index reminder: inject dynamic MEMORY.md content before the user's query. + if !hasMemoryIndexInjected(nRunCtx.AgentInput.Messages) { + indexMsg, err := m.buildMemoryIndexMessage(ctx) + if err != nil { + m.onErr(ctx, OnErrorStageRenderInstruction, err) + } else if !isNilMessage(indexMsg) { + m.sendTopicMemoryEvent(ctx, nRunCtx.AgentInput.Messages, indexMsg) + reminders = append(reminders, indexMsg) + } + } + + // 3) Topic memories: sync mode selects from the original user query. + if !hasTopicMemoryInjected(nRunCtx.AgentInput.Messages) && + m.cfg.Read.Mode == ReadModeSync && m.topicSelectionEnabled() { + memMsg, err := m.selectAndBuildTopicMemoryMessage(ctx, nRunCtx.AgentInput) + if err != nil { + m.onErr(ctx, OnErrorStageTopicSelectionSync, err) + } else if !isNilMessage(memMsg) { + m.sendTopicMemoryEvent(ctx, nRunCtx.AgentInput.Messages, memMsg) + reminders = append(reminders, memMsg) + } + } + + if len(reminders) > 0 { + msgs := insertMessagesBeforeLastUserQuery(nRunCtx.AgentInput.Messages, reminders) + nRunCtx.AgentInput = &adk.TypedAgentInput[M]{Messages: msgs, EnableStreaming: nRunCtx.AgentInput.EnableStreaming} + } + + // 4) Topic memories: async mode starts selection here (cannot use RunLocalValue in BeforeAgent). + if !hasTopicMemoryInjected(nRunCtx.AgentInput.Messages) && + m.cfg.Read.Mode == ReadModeAsync && m.topicSelectionEnabled() { + if existing, _ := ctx.Value(ctxKeySelectionFuture{}).(*selectionFuture); existing == nil { + fut := &selectionFuture{done: make(chan struct{})} + ctx = context.WithValue(ctx, ctxKeySelectionFuture{}, fut) + + // Snapshot current messages for selection; async path is best-effort. + msgSnapshot := append([]M{}, nRunCtx.AgentInput.Messages...) + go func() { + defer close(fut.done) + memMsg, selErr := m.selectAndBuildTopicMemoryMessage(ctx, &adk.TypedAgentInput[M]{Messages: msgSnapshot}) + fut.mu.Lock() + defer fut.mu.Unlock() + if selErr != nil { + fut.err = selErr + return + } + if !isNilMessage(memMsg) { + fut.content = userMessageTextContent(memMsg) + } + }() + } + } + + return ctx, &nRunCtx, nil +} + +func (m *middleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) { + if state == nil { + return ctx, state, nil + } + // Best-effort protection: if automemory content has been injected before and later + // mutated by other components, restore it using the immutable snapshot stored in the future. + if fut, _ := ctx.Value(ctxKeySelectionFuture{}).(*selectionFuture); fut != nil { + fut.mu.Lock() + expected := fut.content + fut.mu.Unlock() + if strings.TrimSpace(expected) != "" { + state = ensureMemoryMsgUnchanged(state, expected) + } + } + if m.cfg.Read.Mode != ReadModeAsync { + return ctx, state, nil + } + if !m.topicSelectionEnabled() { + return ctx, state, nil + } + fut, _ := ctx.Value(ctxKeySelectionFuture{}).(*selectionFuture) + if fut == nil { + return ctx, state, nil + } + + select { + case <-fut.done: + default: + return ctx, state, nil + } + + fut.mu.Lock() + if fut.applied { + fut.mu.Unlock() + return ctx, state, nil + } + content := fut.content + err := fut.err + fut.mu.Unlock() + if err != nil { + m.onErr(ctx, OnErrorStageTopicSelectionAsync, err) + } + + var msgs []M + if strings.TrimSpace(content) != "" { + memMsg := newMemoryMessage[M](content) + m.sendTopicMemoryEvent(ctx, state.Messages, memMsg) + msgs = append(msgs, state.Messages...) + msgs = append(msgs, memMsg) + } else { + msgs = state.Messages + } + + fut.mu.Lock() + fut.applied = true + fut.mu.Unlock() + + return ctx, &adk.TypedChatModelAgentState[M]{Messages: msgs}, nil +} + +type topicSelectionResp struct { + SelectedMemories []string `json:"selected_memories"` +} + +func (m *middleware[M]) renderInstruction(ctx context.Context, baseInstruction string) (string, error) { + memDesc := getDefaultMemoryInstruction() + if m.cfg.GenInstruction != nil { + custom, err := m.cfg.GenInstruction(ctx) + if err != nil { + return "", err + } + if strings.TrimSpace(custom) != "" { + memDesc = custom + "\n\n" + } + } + + return buildSystemMemoryInstruction(baseInstruction, memDesc, m.resolvedMemoryDirectory) +} + +func (m *middleware[M]) buildMemoryIndexMessage(ctx context.Context) (M, error) { + indexPath := filepath.Join(m.resolvedMemoryDirectory, m.cfg.Read.Index.FileName) + indexContent := "" + totalLines := 0 + + fc, err := m.boundedMemoryBackend.Read(ctx, &ReadRequest{FilePath: indexPath}) + if err == nil && fc != nil && !isFileNotFoundContent(fc.Content) { + indexContent = fc.Content + totalLines = strings.Count(indexContent, "\n") + 1 + } + truncatedMemoryIndex, _, truncated := linesOrSizeTrunc(indexContent, m.cfg.Read.Index.MaxLines, m.cfg.Read.Index.MaxBytes) + index := memoryIndexPromptInfo{ + FileName: m.cfg.Read.Index.FileName, + Path: indexPath, + Content: truncatedMemoryIndex, + Empty: strings.TrimSpace(indexContent) == "", + Truncated: truncated, + Lines: totalLines, + IncludeContent: true, + } + return newMemoryIndexMessage[M](buildMemoryIndexReminder(index)), nil +} + +type topicFrontmatter struct { + Name string `yaml:"name"` + Description string `yaml:"description"` + Type string `yaml:"type"` +} + +type topicCandidateBundle struct { + Key string + AbsPath string + RelPath string + Info FileInfo +} + +type topicMemoryPromptInfo struct { + MemoryDirectory string + Path string + Saved string + Content string +} + +func (m *middleware[M]) selectAndBuildTopicMemoryMessage(ctx context.Context, agentIn *adk.TypedAgentInput[M]) (M, error) { + last, ok := m.lastUserMessage(agentIn) + if !ok { + return nil, nil + } + + relToBundle, available, orderedRel, err := m.listTopicCandidates(ctx) + if err != nil || len(orderedRel) == 0 { + return nil, err + } + + topK := m.topicSelectionTopK() + selected, err := m.selectTopicCandidates(ctx, agentIn, userMessageTextContent(last), available, orderedRel, relToBundle) + if err != nil || len(selected) == 0 { + return nil, err + } + + topics := m.renderTopicMemories(ctx, selected, relToBundle, topK) + if len(topics) == 0 { + return nil, nil + } + + return newMemoryMessage[M]("\n" + buildTopicMemoryReminder(topics)), nil +} + +func (m *middleware[M]) listTopicCandidates(ctx context.Context) (map[string]topicCandidateBundle, []string, []string, error) { + candidates, err := m.topicSelectionCandidates(ctx) + if err != nil || len(candidates) == 0 { + return nil, nil, nil, err + } + + relToBundle := make(map[string]topicCandidateBundle, len(candidates)) + available := make([]string, 0, len(candidates)) + orderedRel := make([]string, 0, len(candidates)) + + for _, fi := range candidates { + bundle, manifestLine, ok := m.buildTopicCandidateBundle(ctx, fi) + if !ok { + continue + } + relToBundle[bundle.Key] = bundle + available = append(available, manifestLine) + orderedRel = append(orderedRel, bundle.Key) + } + + return relToBundle, available, orderedRel, nil +} + +func (m *middleware[M]) topicSelectionCandidates(ctx context.Context) ([]topicCandidateBundle, error) { + files, err := m.boundedMemoryBackend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: m.cfg.Read.TopicSelection.CandidateGlob, + Path: m.resolvedMemoryDirectory, + }) + if err != nil { + return nil, err + } + + var candidates []topicCandidateBundle + indexAbs := filepath.Join(m.resolvedMemoryDirectory, m.cfg.Read.Index.FileName) + for _, fi := range files { + if filepath.Clean(fi.Path) == filepath.Clean(indexAbs) { + continue + } + rel, relErr := filepath.Rel(m.resolvedMemoryDirectory, fi.Path) + if relErr != nil { + rel = filepath.Base(fi.Path) + } + rel = filepath.ToSlash(rel) + candidates = append(candidates, topicCandidateBundle{ + Key: rel, + AbsPath: fi.Path, + RelPath: rel, + Info: fi, + }) + } + if len(candidates) == 0 { + return nil, nil + } + + sort.Slice(candidates, func(i, j int) bool { + return parseRFC3339NanoBestEffort(candidates[i].Info.ModifiedAt).After(parseRFC3339NanoBestEffort(candidates[j].Info.ModifiedAt)) + }) + if len(candidates) > m.cfg.Read.TopicSelection.CandidateLimit { + candidates = candidates[:m.cfg.Read.TopicSelection.CandidateLimit] + } + return candidates, nil +} + +func (m *middleware[M]) buildTopicCandidateBundle(ctx context.Context, bundle topicCandidateBundle) (topicCandidateBundle, string, bool) { + preview, err := m.boundedMemoryBackend.Read(ctx, &ReadRequest{ + FilePath: bundle.AbsPath, + Limit: m.cfg.Read.TopicSelection.CandidatePreviewLines, + }) + if err != nil || preview == nil || isFileNotFoundContent(preview.Content) { + return topicCandidateBundle{}, "", false + } + + desc := describeTopicCandidate(preview.Content) + manifestLine := fmt.Sprintf("- %s (saved %s): %s", bundle.Key, bundle.Info.ModifiedAt, desc) + return bundle, manifestLine, true +} + +func (m *middleware[M]) selectTopicCandidates( + ctx context.Context, + agentIn *adk.TypedAgentInput[M], + userQuery string, + available []string, + orderedRel []string, + relToBundle map[string]topicCandidateBundle, +) ([]string, error) { + topK := m.topicSelectionTopK() + + userMsg, err := pyfmt.Fmt(getTopicSelectionUserPrompt(), map[string]any{ + "user_query": userQuery, + "top_k": topK, + "available_memories": strings.Join(available, "\n"), + "tools": strings.Join(collectToolNames(agentIn.Messages), ", "), + }) + if err != nil { + return nil, err + } + + toolInfo := topicSelectionToolInfo() + resp, err := m.topicSelectionModel.Generate( + ctx, + []M{ + makeSystemMsg[M](getTopicSelectionSystemPrompt()), + makeUserMsg[M](userMsg), + }, + makeToolChoiceForced[M](toolInfo.Name), + ) + if err != nil { + return nil, err + } + + valid := make(map[string]struct{}, len(relToBundle)) + for k := range relToBundle { + valid[k] = struct{}{} + } + selected, err := parseTopicSelectionFromToolCall(resp, valid) + if err != nil { + return nil, err + } + if len(selected) > topK { + return selected[:topK], nil + } + return selected, nil +} + +func (m *middleware[M]) renderTopicMemories( + ctx context.Context, + selected []string, + relToBundle map[string]topicCandidateBundle, + topK int, +) []topicMemoryPromptInfo { + capHint := topK + if capHint > len(selected) { + capHint = len(selected) + } + rendered := make([]topicMemoryPromptInfo, 0, capHint) + totalBytes := 0 + maxTotalBytes := m.cfg.Read.TopicSelection.MaxTotalBytes + for _, rel := range selected { + if len(rendered) >= topK { + break + } + bundle, ok := relToBundle[rel] + if !ok { + continue + } + topic, ok := m.renderTopicMemory(ctx, bundle) + if !ok { + continue + } + topicBytes := len(topic.Content) + len(topic.MemoryDirectory) + len(topic.Path) + if maxTotalBytes > 0 && totalBytes+topicBytes > maxTotalBytes { + if len(rendered) == 0 { + if len(topic.Content) > maxTotalBytes { + topic.Content = topic.Content[:maxTotalBytes] + } + rendered = append(rendered, topic) + } + break + } + rendered = append(rendered, topic) + totalBytes += topicBytes + } + return rendered +} + +func (m *middleware[M]) renderTopicMemory(ctx context.Context, bundle topicCandidateBundle) (topicMemoryPromptInfo, bool) { + full, err := m.boundedMemoryBackend.Read(ctx, &ReadRequest{FilePath: bundle.AbsPath}) + if err != nil || full == nil || isFileNotFoundContent(full.Content) { + return topicMemoryPromptInfo{}, false + } + + content, truncReason, truncated := linesOrSizeTrunc(full.Content, m.cfg.Read.TopicSelection.MaxLines, m.cfg.Read.TopicSelection.MaxBytes) + if truncated { + truncNotify, err := pyfmt.Fmt(getTopicMemoryTruncNotify(), map[string]any{ + "reason": truncReason, + "abs_path": bundle.AbsPath, + }) + if err == nil { + content += truncNotify + } + } + + return topicMemoryPromptInfo{ + MemoryDirectory: m.resolvedMemoryDirectory, + Path: bundle.RelPath, + Saved: bundle.Info.ModifiedAt, + Content: content, + }, true +} + +func (m *middleware[M]) AfterAgent(ctx context.Context, state *adk.TypedChatModelAgentState[M]) (context.Context, error) { + if m.cfg == nil || m.cfg.Write == nil || m.cfg.Write.Mode == WriteModeDisabled { + return ctx, nil + } + if m.cfg.Write.Model == nil || m.extractionHandler == nil { + return ctx, nil + } + if state == nil || len(state.Messages) == 0 { + return ctx, nil + } + + sessionID, err := m.resolveSessionID(ctx, state) + if err != nil { + m.onErr(ctx, OnErrorStageResolveSessionID, err) + return ctx, nil + } + coordKey := m.coordinatorKey(sessionID) + + cursor := getWriteCursorFromMessages(state.Messages) + if coordKey != "" { + if remoteCursor, ok, err := getCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey); err == nil && ok && remoteCursor > cursor { + cursor = remoteCursor + state = markWriteCursor(state, cursor) + } + } + if cursor >= len(state.Messages) { + return ctx, nil + } + + // Skip background extraction if the main agent already wrote memory files in this range. + if hasMemoryWritesSince(state.Messages, cursor, m.resolvedMemoryDirectory) { + end := len(state.Messages) + if coordKey != "" { + _ = setCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey, end) + } + state = markWriteCursor(state, end) + return ctx, nil + } + + if countModelVisibleMessages(state.Messages[cursor:]) == 0 { + end := len(state.Messages) + if coordKey != "" { + _ = setCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey, end) + } + state = markWriteCursor(state, end) + return ctx, nil + } + + switch m.cfg.Write.Mode { + case WriteModeDisabled: + // do nothing + return ctx, nil + + case WriteModeSync: + end := len(state.Messages) + if err := m.runMemoryExtractionAgent(ctx, state.Messages, cursor, state.ToolInfos); err != nil { + m.onErr(ctx, OnErrorStageMemoryWriteSync, err) + return ctx, nil + } + if coordKey != "" { + _ = setCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey, end) + } + state = markWriteCursor(state, end) + return ctx, nil + + case WriteModeAsync: + if coordKey == "" { + if err := m.runMemoryExtractionAgent(ctx, state.Messages, cursor, state.ToolInfos); err != nil { + m.onErr(ctx, OnErrorStageMemoryWriteSync, err) + return ctx, nil + } + state = markWriteCursor(state, len(state.Messages)) + return ctx, nil + } + snap, err := buildPendingSnapshot(state.Messages, cursor, state.ToolInfos) + if err != nil { + m.onErr(ctx, OnErrorStageSnapshotMarshal, err) + return ctx, nil + } + unlock, ok, err := m.coordination.Coordinator.AcquireLock(ctx, coordKey, m.coordination.LockTTL) + if err != nil { + m.onErr(ctx, OnErrorStageAcquireExtractionLock, err) + return ctx, nil + } + if !ok { + if err := setCoordinatorPendingSnapshot(ctx, m.coordination.Coordinator, coordKey, snap, m.coordination.LockTTL); err != nil { + m.onErr(ctx, OnErrorStageStashPendingSnapshot, err) + } + return ctx, nil + } + go m.runExtractionDrain(ctx, coordKey, unlock, snap) + return ctx, nil + + default: + return ctx, nil + } +} + +func (m *middleware[M]) runExtractionDrain(ctx context.Context, coordKey string, unlock func(context.Context) error, initial *PendingSnapshot) { + defer func() { + if unlock == nil { + return + } + if err := unlock(ctx); err != nil { + m.onErr(ctx, OnErrorStageReleaseExtractionLock, err) + } + }() + + current := initial + for current != nil { + msgs, cursor, toolInfos, err := decodePendingSnapshot[M](current) + if err != nil { + m.onErr(ctx, OnErrorStageDecodePendingSnapshot, err) + } else if err := m.runMemoryExtractionAgent(ctx, msgs, cursor, toolInfos); err != nil { + m.onErr(ctx, OnErrorStageMemoryWriteAsync, err) + } else { + _ = setCoordinatorCursor(ctx, m.coordination.Coordinator, coordKey, len(msgs)) + } + + next, loadErr := popCoordinatorPendingSnapshot(ctx, m.coordination.Coordinator, coordKey) + if loadErr != nil { + m.onErr(ctx, OnErrorStageLoadPendingSnapshot, loadErr) + return + } + current = next + } +} + +func (m *middleware[M]) newExtractionAgent(ctx context.Context, toolInfos []*schema.ToolInfo) (*adk.TypedChatModelAgent[M], error) { + if m.cfg == nil || m.cfg.Write == nil || m.cfg.Write.Model == nil { + return nil, fmt.Errorf("auto memory extraction agent init failed: missing write model") + } + if m.extractionHandler == nil { + return nil, fmt.Errorf("auto memory extraction agent init failed: missing extraction handler") + } + + agent, err := adk.NewTypedChatModelAgent[M](ctx, &adk.TypedChatModelAgentConfig[M]{ + Name: "automemory_extractor", + Model: m.cfg.Write.Model, + Handlers: []adk.TypedChatModelAgentMiddleware[M]{ + m.extractionHandler, // fs middleware + &toolInfoOverrideMiddleware[M]{toolInfos: toolInfos}, // tool info override, for prefix cache + }, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + UnknownToolsHandler: func(ctx context.Context, name, input string) (string, error) { + return "This tool is not allowed to be called. Please follow user prompt to proceed.", nil + }, + }, + EmitInternalEvents: false, + }, + MaxIterations: m.cfg.Write.MaxTurns, + }) + if err != nil { + return nil, fmt.Errorf("auto memory extraction agent init failed: %w", err) + } + return agent, nil +} + +func (m *middleware[M]) runMemoryExtractionAgent(ctx context.Context, snapshot []M, cursor int, toolInfos []*schema.ToolInfo) error { + if len(snapshot) == 0 || cursor >= len(snapshot) { + return nil + } + manifest, err := m.buildMemoryManifest(ctx) + if err != nil { + return err + } + newMessageCount := countModelVisibleMessagesSince(snapshot, cursor) + savePolicy, err := m.extractSavePolicyInstruction(ctx) + if err != nil { + return err + } + userPrompt := buildExtractAutoOnlyPrompt(m.extractionMemoryDirectoryPrompt(), newMessageCount, manifest, savePolicy) + msgs := append(append([]M{}, snapshot...), makeUserMsg[M](userPrompt)) + extractionAgent, err := m.newExtractionAgent(ctx, toolInfos) + if err != nil { + return err + } + + iter := extractionAgent.Run(ctx, &adk.TypedAgentInput[M]{ + Messages: msgs, + EnableStreaming: true, + }) + + if m.cfg != nil && m.cfg.Write != nil && m.cfg.Write.HandleExtractionIterator != nil { + return m.cfg.Write.HandleExtractionIterator(ctx, iter) + } + + for { + ev, ok := iter.Next() + if !ok { + return nil + } + if ev == nil { + continue + } + if ev.Err != nil { + return ev.Err + } + } +} + +func (m *middleware[M]) extractSavePolicyInstruction(ctx context.Context) (string, error) { + if m.cfg == nil || m.cfg.Write == nil || m.cfg.Write.GenInstruction == nil { + return "", nil + } + custom, err := m.cfg.Write.GenInstruction(ctx) + if err != nil { + return "", err + } + return strings.TrimSpace(custom), nil +} + +func (m *middleware[M]) extractionMemoryDirectoryPrompt() string { + index := &memoryIndexPromptInfo{ + FileName: m.cfg.Read.Index.FileName, + Path: filepath.Join(m.resolvedMemoryDirectory, m.cfg.Read.Index.FileName), + } + return buildMemoryDirectoryManifest(m.resolvedMemoryDirectory, index) +} + +func (m *middleware[M]) buildMemoryManifest(ctx context.Context) (string, error) { + files, err := m.boundedMemoryBackend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: CandidateGlobPattern, + Path: m.resolvedMemoryDirectory, + }) + if err != nil { + return "", err + } + manifest := memoryManifestPromptInfo{Directory: m.resolvedMemoryDirectory} + indexAbs := filepath.Join(m.resolvedMemoryDirectory, m.cfg.Read.Index.FileName) + for _, fi := range files { + rel, relErr := filepath.Rel(m.resolvedMemoryDirectory, fi.Path) + if relErr != nil { + rel = filepath.Base(fi.Path) + } + rel = filepath.ToSlash(rel) + if filepath.Clean(fi.Path) == filepath.Clean(indexAbs) { + rel = m.cfg.Read.Index.FileName + } + desc := "" + preview, rerr := m.boundedMemoryBackend.Read(ctx, &ReadRequest{FilePath: fi.Path, Limit: defaultCandidatePreviewLine}) + if rerr == nil && preview != nil && !isFileNotFoundContent(preview.Content) { + if fm, ok := parseFrontmatter(preview.Content); ok { + desc = strings.TrimSpace(fm.Description) + } + } + manifest.Files = append(manifest.Files, memoryManifestFilePromptInfo{ + MemoryPath: rel, + AbsPath: fi.Path, + Saved: fi.ModifiedAt, + Description: desc, + }) + } + return buildExtractionMemoryManifest(manifest), nil +} + +type toolInfoOverrideMiddleware[M adk.MessageType] struct { + adk.TypedBaseChatModelAgentMiddleware[M] + + toolInfos []*schema.ToolInfo +} + +func (t *toolInfoOverrideMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) ( + context.Context, *adk.TypedChatModelAgentState[M], error) { + + toolNameMapping := make(map[string]struct{}, len(t.toolInfos)) + for _, tool := range t.toolInfos { + toolNameMapping[tool.Name] = struct{}{} + } + + overrideTools := append([]*schema.ToolInfo{}, t.toolInfos...) + for _, tool := range state.ToolInfos { + if _, ok := toolNameMapping[tool.Name]; !ok { + overrideTools = append(overrideTools, tool) + } + } + state.ToolInfos = overrideTools + + return ctx, state, nil +} + +type modelWithTools[M adk.MessageType] struct { + base model.BaseModel[M] + tools []*schema.ToolInfo +} + +func (m *modelWithTools[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + newOpts := make([]model.Option, len(opts)+1) + copy(newOpts, opts) + newOpts[len(opts)] = model.WithTools(m.tools) + return m.base.Generate(ctx, input, newOpts...) +} + +func (m *modelWithTools[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + newOpts := make([]model.Option, len(opts)+1) + copy(newOpts, opts) + newOpts[len(opts)] = model.WithTools(m.tools) + return m.base.Stream(ctx, input, newOpts...) +} diff --git a/adk/middlewares/automemory/automemory_test.go b/adk/middlewares/automemory/automemory_test.go new file mode 100644 index 000000000..38f9584d5 --- /dev/null +++ b/adk/middlewares/automemory/automemory_test.go @@ -0,0 +1,1532 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type fixedModel struct { + out string +} + +func (m *fixedModel) Generate(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "select-fixed", + Type: "function", + Function: schema.FunctionCall{ + Name: topicSelectionToolName, + Arguments: m.out, + }, + }, + }), nil +} + +func (m *fixedModel) Stream(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, _ := m.Generate(ctx, input) + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *fixedModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func requireMemoryIndexMessage(t *testing.T, msg *schema.Message, contains ...string) { + t.Helper() + require.True(t, isMemoryIndexMessage(msg)) + require.NotNil(t, msg.Extra) + require.NotNil(t, msg.Extra[memoryExtraKey]) + require.Contains(t, msg.Content, "") + require.Contains(t, msg.Content, "") + require.True(t, + strings.Contains(msg.Content, "# Memory Index") || strings.Contains(msg.Content, "# 记忆索引文件"), + "memory index reminder should contain an index title", + ) + require.NotContains(t, msg.Content, "") + require.NotContains(t, msg.Content, "") + require.NotContains(t, msg.Content, "") + require.NotContains(t, msg.Content, "") + require.NotContains(t, msg.Content, "### 1. Name:") + require.NotContains(t, msg.Content, "#### Index file content:") + for _, s := range contains { + require.Contains(t, msg.Content, s) + } +} + +func requireTopicMemoryMessage(t *testing.T, msg *schema.Message, contains ...string) { + t.Helper() + require.True(t, isTopicMemoryMessage(msg)) + require.NotNil(t, msg.Extra) + require.NotNil(t, msg.Extra[memoryExtraKey]) + require.Contains(t, msg.Content, "") + require.Contains(t, msg.Content, "") + require.Contains(t, msg.Content, "") + require.NotContains(t, msg.Content, "") + require.NotContains(t, msg.Content, "") + require.Contains(t, msg.Content, "") + for _, s := range contains { + require.Contains(t, msg.Content, s) + } +} + +func requireWriteCursor(t *testing.T, msgs []*schema.Message, cursor int) { + t.Helper() + for _, msg := range msgs { + if msg == nil || msg.Extra == nil { + continue + } + meta, ok := msg.Extra[memoryExtraKey].(*memoryExtra) + if ok && meta != nil && meta.Type == "write_cursor" { + require.EqualValues(t, cursor, meta.Cursor) + return + } + } + require.Fail(t, "write cursor not found") +} + +func countMemoryIndexMessages(msgs []*schema.Message) int { + count := 0 + for _, msg := range msgs { + if isMemoryIndexMessage(msg) { + count++ + } + } + return count +} + +func countTopicMemoryMessages(msgs []*schema.Message) int { + count := 0 + for _, msg := range msgs { + if isTopicMemoryMessage(msg) { + count++ + } + } + return count +} + +func TestMiddleware_IndexInjection_Empty(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + // Model nil => topic selection disabled. + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Contains(t, out.Instruction, "# Auto memory") + require.Contains(t, out.Instruction, "## Memory directory") + require.Contains(t, out.Instruction, "Path: /mem") + require.NotContains(t, out.Instruction, "Index file path: /mem/MEMORY.md") + require.NotContains(t, out.Instruction, "#### Index file content: MEMORY.md") + require.NotContains(t, out.Instruction, "Rules:") + require.Len(t, out.AgentInput.Messages, 2) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md", "currently empty") + require.Contains(t, out.AgentInput.Messages[1].Content, "hi") +} + +func TestMiddleware_IndexInjection_ChineseInstruction(t *testing.T) { + require.NoError(t, adk.SetLanguage(adk.LanguageChinese)) + defer func() { + require.NoError(t, adk.SetLanguage(adk.LanguageEnglish)) + }() + + ctx := context.Background() + b := NewInMemoryBackend() + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Contains(t, out.Instruction, "# 自动记忆") + require.NotContains(t, out.Instruction, "你的 MEMORY.md 当前为空") + require.Len(t, out.AgentInput.Messages, 2) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "# 记忆索引文件", "文件 /mem/MEMORY.md", "内容为空") + require.Contains(t, out.AgentInput.Messages[1].Content, "hi") +} + +func TestMiddleware_IndexInjection_CustomInstructionKeepsDirectoryManifest(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + custom := "custom memory header" + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + GenInstruction: func(ctx context.Context) (string, error) { + return custom, nil + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Contains(t, out.Instruction, "custom memory header") + require.Contains(t, out.Instruction, "## Memory directory") + require.Contains(t, out.Instruction, "Path: /mem") + require.NotContains(t, out.Instruction, "Index file path") + require.Len(t, out.AgentInput.Messages, 2) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + require.Contains(t, out.AgentInput.Messages[1].Content, "hi") +} + +func TestMiddleware_IndexInjection_CustomInstructionErrorReportsRenderStage(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + var stages []ErrorStage + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + GenInstruction: func(ctx context.Context) (string, error) { + return "", fmt.Errorf("custom instruction failed") + }, + OnError: func(ctx context.Context, stage ErrorStage, err error) { + stages = append(stages, stage) + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Equal(t, "base", out.Instruction) + require.Equal(t, []ErrorStage{OnErrorStageRenderInstruction}, stages) +} + +func TestNew_DoesNotMutateConfig(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + + cfgNilNested := &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &fixedModel{out: `{"selected_memories":["debugging.md"]}`}, + } + _, err := New(ctx, cfgNilNested) + require.NoError(t, err) + require.Nil(t, cfgNilNested.Read) + require.Nil(t, cfgNilNested.Write) + require.Nil(t, cfgNilNested.Coordination) + + cfgExplicitNested := &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &fixedModel{out: `{"selected_memories":["debugging.md"]}`}, + Read: &ReadConfig[*schema.Message]{}, + Write: &WriteConfig[*schema.Message]{}, + Coordination: &CoordinationConfig[*schema.Message]{}, + } + _, err = New(ctx, cfgExplicitNested) + require.NoError(t, err) + require.Empty(t, cfgExplicitNested.Read.Mode) + require.Nil(t, cfgExplicitNested.Read.Model) + require.Nil(t, cfgExplicitNested.Read.Index) + require.Nil(t, cfgExplicitNested.Read.TopicSelection) + require.Empty(t, cfgExplicitNested.Write.Mode) + require.Nil(t, cfgExplicitNested.Write.Model) + require.Zero(t, cfgExplicitNested.Write.MaxTurns) + require.Nil(t, cfgExplicitNested.Coordination.Coordinator) + require.Zero(t, cfgExplicitNested.Coordination.LockTTL) +} + +func TestMiddleware_TopicSelection_InsertsMemoryMessage(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md) - notes\n", now) + b.put("/mem/debugging.md", "---\nname: Debugging\ndescription: build and test commands\ntype: project\n---\n\n# Debugging\npnpm test\n", now) + b.put("/mem/other.md", "---\nname: Other\ndescription: unrelated\ntype: misc\n---\n", now.Add(-time.Hour)) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &fixedModel{out: `{"selected_memories":["debugging.md"]}`}, + }) + require.NoError(t, err) + + in := &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to run tests?")}} + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: in, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.NotNil(t, out.AgentInput) + require.Len(t, out.AgentInput.Messages, 3) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + requireTopicMemoryMessage(t, out.AgentInput.Messages[1], "Contents of /mem/debugging.md") + require.Equal(t, schema.User, out.AgentInput.Messages[2].Role) + require.Contains(t, out.AgentInput.Messages[2].Content, "How to run tests?") +} + +func TestMiddleware_MemoryDirectory_IndexAndTopicSelection(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [prefs.md](prefs.md) - user preferences\n- [debugging.md](debugging.md) - project debugging\n", now) + b.put("/mem/prefs.md", "---\ndescription: editor preferences\n---\n\nUse concise answers.\n", now) + b.put("/mem/debugging.md", "---\ndescription: test commands\n---\n\nRun go test ./...\n", now) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &fixedModel{out: `{"selected_memories":["debugging.md"]}`}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How should I run tests?")}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Contains(t, out.Instruction, "## Memory directory") + require.Contains(t, out.Instruction, "Path: /mem") + require.NotContains(t, out.Instruction, "Index file path: /mem/MEMORY.md") + require.NotContains(t, out.Instruction, "#### Index file content: MEMORY.md") + require.Len(t, out.AgentInput.Messages, 3) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], + "Contents of /mem/MEMORY.md", + "- [prefs.md](prefs.md) - user preferences", + "- [debugging.md](debugging.md) - project debugging", + ) + indexReminder := out.AgentInput.Messages[0].Content + userIndexPos := strings.Index(indexReminder, "- [prefs.md](prefs.md) - user preferences") + projectIndexPos := strings.Index(indexReminder, "- [debugging.md](debugging.md) - project debugging") + require.True(t, userIndexPos >= 0 && projectIndexPos > userIndexPos) + requireTopicMemoryMessage(t, out.AgentInput.Messages[1], "Contents of /mem/debugging.md", "Run go test ./...") + require.NotContains(t, out.AgentInput.Messages[1].Content, "Use concise answers.") + require.Contains(t, out.AgentInput.Messages[2].Content, "How should I run tests?") +} + +func TestMiddleware_TopicSelection_AsyncInjectsInBeforeModel(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md) - notes\n", now) + b.put("/mem/debugging.md", "---\nname: Debugging\ndescription: build and test commands\ntype: project\n---\n\n# Debugging\npnpm test\n", now) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &fixedModel{out: `{"selected_memories":["debugging.md"]}`}, + Read: &ReadConfig[*schema.Message]{Mode: ReadModeAsync}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to run tests?")}}, + } + ctx2, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Len(t, out.AgentInput.Messages, 2) // async doesn't inject topic memory here + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + require.Contains(t, out.AgentInput.Messages[1].Content, "How to run tests?") + + st := &adk.ChatModelAgentState{Messages: []adk.Message{schema.UserMessage("How to run tests?")}} + + require.Eventually(t, func() bool { + _, next, err := mw.BeforeModelRewriteState(ctx2, st, nil) + require.NoError(t, err) + st = next + last := st.Messages[len(st.Messages)-1] + return len(st.Messages) == 2 && last.Extra != nil && last.Extra["__eino_automemory__"] != nil + }, 2*time.Second, 10*time.Millisecond) +} + +type panicModel struct{} + +func (m *panicModel) Generate(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + panic("should not call model") +} + +func (m *panicModel) Stream(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("should not call model") +} + +func (m *panicModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +type toolCallSelectionModel struct { + calls int32 +} + +func (m *toolCallSelectionModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m.calls, 1) + return schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "select-1", + Type: "function", + Function: schema.FunctionCall{ + Name: topicSelectionToolName, + Arguments: `{"selected_memories":["debugging.md","hallucinated.md"]}`, + }, + }, + }), nil +} + +func (m *toolCallSelectionModel) Stream(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *toolCallSelectionModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +type extractionModel struct { + mu sync.Mutex + promptSeen []string + boundToolCalls [][]string + topicPath string + indexPath string + blockFirstRun chan struct{} + firstRunStarted chan struct{} + blockedOnce uint32 // atomic (0/1) + generateCallings int32 +} + +type countingBackend struct { + *InMemoryBackend + writeCalls int32 + globInfoCalls int32 + mu sync.Mutex + paths []string +} + +type outOfBoundsCandidateBackend struct { + outsideReadCalled int32 +} + +func (b *outOfBoundsCandidateBackend) Read(_ context.Context, req *ReadRequest) (*FileContent, error) { + if req == nil { + return nil, fmt.Errorf("read: invalid request") + } + if filepath.Clean(req.FilePath) == filepath.Clean("/outside/secret.md") { + atomic.StoreInt32(&b.outsideReadCalled, 1) + return &FileContent{Content: "secret"}, nil + } + return nil, fmt.Errorf("file not found: %s", req.FilePath) +} + +func (b *outOfBoundsCandidateBackend) GlobInfo(_ context.Context, req *GlobInfoRequest) ([]FileInfo, error) { + if req == nil { + return nil, fmt.Errorf("glob: invalid request") + } + return []FileInfo{{ + Path: "/outside/secret.md", + ModifiedAt: time.Now().Format(time.RFC3339Nano), + }}, nil +} + +func (b *outOfBoundsCandidateBackend) Write(context.Context, *WriteRequest) error { + return nil +} + +func (b *outOfBoundsCandidateBackend) Edit(context.Context, *EditRequest) error { + return nil +} + +func (b *countingBackend) Write(ctx context.Context, req *WriteRequest) error { + atomic.AddInt32(&b.writeCalls, 1) + b.mu.Lock() + b.paths = append(b.paths, req.FilePath) + b.mu.Unlock() + return b.InMemoryBackend.Write(ctx, req) +} + +func (b *countingBackend) GlobInfo(ctx context.Context, req *GlobInfoRequest) ([]FileInfo, error) { + atomic.AddInt32(&b.globInfoCalls, 1) + return b.InMemoryBackend.GlobInfo(ctx, req) +} + +func (m *extractionModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m.generateCallings, 1) + promptIdx := findExtractionPromptIndex(input) + if promptIdx < 0 { + return nil, fmt.Errorf("missing extraction prompt") + } + + m.mu.Lock() + m.promptSeen = append(m.promptSeen, input[promptIdx].Content) + m.mu.Unlock() + + if hasToolMessageAfter(input, promptIdx) { + return schema.AssistantMessage("done", nil), nil + } + + if m.blockFirstRun != nil && atomic.SwapUint32(&m.blockedOnce, 1) == 0 { + if m.firstRunStarted != nil { + close(m.firstRunStarted) + } + <-m.blockFirstRun + } + + payload := lastBusinessUserBeforePrompt(input, promptIdx) + topicPath := m.topicPath + if topicPath == "" { + topicPath = "topic.md" + } + indexPath := m.indexPath + if indexPath == "" { + indexPath = "MEMORY.md" + } + return schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "write-topic", + Type: "function", + Function: schema.FunctionCall{ + Name: "write_file", + Arguments: fmt.Sprintf(`{"file_path":%q,"content":%q}`, topicPath, payload), + }, + }, + { + ID: "write-index", + Type: "function", + Function: schema.FunctionCall{ + Name: "write_file", + Arguments: fmt.Sprintf(`{"file_path":%q,"content":"- [topic.md](topic.md)\n"}`, indexPath), + }, + }, + }), nil +} + +func (m *extractionModel) Stream(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *extractionModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + names := make([]string, 0, len(tools)) + for _, ti := range tools { + if ti == nil { + continue + } + names = append(names, ti.Name) + } + m.mu.Lock() + m.boundToolCalls = append(m.boundToolCalls, names) + m.mu.Unlock() + return m, nil +} + +func findExtractionPromptIndex(input []*schema.Message) int { + for i := len(input) - 1; i >= 0; i-- { + if input[i] != nil && input[i].Role == schema.User && + (strings.Contains(input[i].Content, "memory extraction subagent") || strings.Contains(input[i].Content, "记忆提取子智能体")) { + return i + } + } + return -1 +} + +func hasToolMessageAfter(input []*schema.Message, idx int) bool { + for i := idx + 1; i < len(input); i++ { + if input[i] != nil && input[i].Role == schema.Tool { + switch input[i].ToolName { + case "read_file", "glob", "write_file", "edit_file": + return true + default: + } + } + } + return false +} + +func lastBusinessUserBeforePrompt(input []*schema.Message, promptIdx int) string { + for i := promptIdx - 1; i >= 0; i-- { + if input[i] == nil || input[i].Role != schema.User { + continue + } + if strings.Contains(input[i].Content, "") { + continue + } + return input[i].Content + } + return "unknown" +} + +func TestMiddleware_TopicSelection_SmallCandidateSetUsesModel(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md)\n- [patterns.md](patterns.md)\n", now) + b.put("/mem/debugging.md", "---\ndescription: debug notes\n---\nbody\n", now) + b.put("/mem/patterns.md", "---\ndescription: patterns\n---\nbody\n", now) + model := &toolCallSelectionModel{} + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: model, + Read: &ReadConfig[*schema.Message]{ + Mode: ReadModeSync, + TopicSelection: &TopicSelectionConfig{ + TopK: 5, + }, + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to run tests?")}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Equal(t, int32(1), atomic.LoadInt32(&model.calls)) + require.Len(t, out.AgentInput.Messages, 3) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + requireTopicMemoryMessage(t, out.AgentInput.Messages[1], "debugging.md") + require.NotContains(t, out.AgentInput.Messages[1].Content, "patterns.md") + require.Contains(t, out.AgentInput.Messages[2].Content, "How to run tests?") +} + +func TestMiddleware_TopicSelection_DisabledSkipsSelectionAndReminder(t *testing.T) { + for _, mode := range []ReadMode{ReadModeSync, ReadModeAsync} { + t.Run(string(mode), func(t *testing.T) { + ctx := context.Background() + b := &countingBackend{InMemoryBackend: NewInMemoryBackend()} + now := time.Now() + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md)\n", now) + b.put("/mem/debugging.md", "---\ndescription: debug notes\n---\nbody\n", now) + + selModel := &toolCallSelectionModel{} + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: selModel, + Read: &ReadConfig[*schema.Message]{ + Mode: mode, + TopicSelection: &TopicSelectionConfig{ + Enable: boolPtr(false), + TopK: 1, + }, + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to debug?")}}, + } + ctx2, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Len(t, out.AgentInput.Messages, 2) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + require.Contains(t, out.AgentInput.Messages[1].Content, "How to debug?") + require.Equal(t, 0, countTopicMemoryMessages(out.AgentInput.Messages)) + require.EqualValues(t, 0, atomic.LoadInt32(&selModel.calls)) + require.EqualValues(t, 0, atomic.LoadInt32(&b.globInfoCalls)) + + if mode == ReadModeAsync { + st := &adk.ChatModelAgentState{Messages: []adk.Message{schema.UserMessage("How to debug?")}} + _, next, err := mw.BeforeModelRewriteState(ctx2, st, nil) + require.NoError(t, err) + require.Len(t, next.Messages, 1) + require.Equal(t, 0, countTopicMemoryMessages(next.Messages)) + require.EqualValues(t, 0, atomic.LoadInt32(&selModel.calls)) + require.EqualValues(t, 0, atomic.LoadInt32(&b.globInfoCalls)) + } + }) + } +} + +func TestMiddleware_AfterAgent_SyncExtractionWritesMemoryFiles(t *testing.T) { + ctx := context.Background() + b := &countingBackend{InMemoryBackend: NewInMemoryBackend()} + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + extModel := &extractionModel{} + var onErrStages []ErrorStage + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeSync, + Model: extModel, + }, + OnError: func(ctx context.Context, stage ErrorStage, err error) { + onErrStages = append(onErrStages, stage) + }, + }) + require.NoError(t, err) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember alpha"), + schema.AssistantMessage("ack", nil), + }, + } + + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{ + Messages: state.Messages, + ToolInfos: []*schema.ToolInfo{ + {Name: "tool_b"}, + {Name: "tool_a"}, + }, + }) + require.NoError(t, err) + require.Empty(t, onErrStages) + require.Equal(t, len(state.Messages), getWriteCursorFromMessages(state.Messages)) + require.GreaterOrEqual(t, atomic.LoadInt32(&extModel.generateCallings), int32(1)) + require.GreaterOrEqual(t, atomic.LoadInt32(&b.writeCalls), int32(1)) + b.mu.Lock() + paths := append([]string(nil), b.paths...) + b.mu.Unlock() + require.NotEmpty(t, paths) + require.Contains(t, paths, "/mem/topic.md") + require.Contains(t, paths, "/mem/MEMORY.md") + + mem, err := b.Read(ctx, &ReadRequest{FilePath: "/mem/MEMORY.md"}) + require.NoError(t, err) + require.Contains(t, mem.Content, "topic.md") + + topic, err := b.Read(ctx, &ReadRequest{FilePath: "/mem/topic.md"}) + require.NoError(t, err) + require.Equal(t, "remember alpha", topic.Content) + + extModel.mu.Lock() + defer extModel.mu.Unlock() + require.NotEmpty(t, extModel.promptSeen) + require.Contains(t, extModel.promptSeen[0], "memory extraction subagent") + require.Contains(t, extModel.promptSeen[0], "## Memory directory") + require.Contains(t, extModel.promptSeen[0], "Path: /mem") +} + +func TestMiddleware_AfterAgent_SyncExtraction_CustomWriteInstruction(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + extModel := &extractionModel{} + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeSync, + Model: extModel, + GenInstruction: func(ctx context.Context) (string, error) { + return "## Custom save policy\n- Save only explicitly requested memories\n- Prefer updating user_preferences.md for user preference changes\n- Do not save temporary debugging notes", nil + }, + }, + }) + require.NoError(t, err) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember beta"), + schema.AssistantMessage("ack", nil), + }, + } + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{Messages: state.Messages}) + require.NoError(t, err) + + extModel.mu.Lock() + defer extModel.mu.Unlock() + require.NotEmpty(t, extModel.promptSeen) + prompt := extModel.promptSeen[0] + require.Contains(t, prompt, "## Custom save policy") + require.Contains(t, prompt, "- Save only explicitly requested memories") + require.Contains(t, prompt, "- Prefer updating user_preferences.md for user preference changes") + require.Contains(t, prompt, "- Do not save temporary debugging notes") + require.NotContains(t, prompt, "## What to save") + require.NotContains(t, prompt, "- Stable patterns and conventions confirmed across multiple interactions") + require.NotContains(t, prompt, "## What NOT to save") + require.NotContains(t, prompt, "- Session-specific temporary state or current task details") + require.Contains(t, prompt, "## How to save memories") +} + +func TestMiddleware_AfterAgent_SyncExtractionWritesMemoryDirectory(t *testing.T) { + ctx := context.Background() + b := &countingBackend{InMemoryBackend: NewInMemoryBackend()} + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + extModel := &extractionModel{ + topicPath: "topic.md", + indexPath: "MEMORY.md", + } + var onErrStages []ErrorStage + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeSync, + Model: extModel, + }, + OnError: func(ctx context.Context, stage ErrorStage, err error) { + onErrStages = append(onErrStages, stage) + }, + }) + require.NoError(t, err) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember project convention"), + schema.AssistantMessage("ack", nil), + }, + } + + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{ + Messages: state.Messages, + }) + require.NoError(t, err) + require.Empty(t, onErrStages) + + topic, err := b.Read(ctx, &ReadRequest{FilePath: "/mem/topic.md"}) + require.NoError(t, err) + require.Equal(t, "remember project convention", topic.Content) + + b.mu.Lock() + paths := append([]string(nil), b.paths...) + b.mu.Unlock() + require.Contains(t, paths, "/mem/topic.md") + require.Contains(t, paths, "/mem/MEMORY.md") +} + +func TestMiddleware_AfterAgent_SyncExtraction_IteratorHandlerCanDrain(t *testing.T) { + ctx := context.Background() + b := &countingBackend{InMemoryBackend: NewInMemoryBackend()} + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + extModel := &extractionModel{} + var seen int32 + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeSync, + Model: extModel, + HandleExtractionIterator: func(ctx context.Context, iter *adk.AsyncIterator[*adk.AgentEvent]) error { + for { + ev, ok := iter.Next() + if !ok { + return nil + } + if ev == nil { + continue + } + atomic.AddInt32(&seen, 1) + if ev.Err != nil { + return ev.Err + } + } + }, + }, + }) + require.NoError(t, err) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember handler"), + schema.AssistantMessage("ack", nil), + }, + } + + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{ + Messages: state.Messages, + ToolInfos: []*schema.ToolInfo{ + {Name: "tool_1"}, + }, + }) + require.NoError(t, err) + require.Greater(t, atomic.LoadInt32(&seen), int32(0)) + + // Still writes memory files as usual (handler only changes event draining). + _, err = b.Read(ctx, &ReadRequest{FilePath: "/mem/topic.md"}) + require.NoError(t, err) +} + +func TestMiddleware_AfterAgent_SkipsExtractionWhenMainAgentAlreadyWroteMemory(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + extModel := &extractionModel{} + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeSync, + Model: extModel, + }, + }) + require.NoError(t, err) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember beta"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call-1", + Type: "function", + Function: schema.FunctionCall{ + Name: "write_file", + Arguments: `{"file_path":"/mem/topic.md","content":"written by main agent"}`, + }, + }, + }), + schema.ToolMessage("ok", "call-1", schema.WithToolName("write_file")), + }, + } + + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{Messages: state.Messages}) + require.NoError(t, err) + require.Equal(t, len(state.Messages), getWriteCursorFromMessages(state.Messages)) + require.EqualValues(t, 0, atomic.LoadInt32(&extModel.generateCallings)) + + _, err = b.Read(ctx, &ReadRequest{FilePath: "/mem/topic.md"}) + require.Error(t, err) +} + +func TestMiddleware_AfterAgent_AsyncExtractionKeepsLatestPendingSnapshot(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + blockCh := make(chan struct{}) + startedCh := make(chan struct{}) + extModel := &extractionModel{ + blockFirstRun: blockCh, + firstRunStarted: startedCh, + } + coord := &CoordinationConfig[*schema.Message]{ + SessionID: "session-1", + Coordinator: NewLocalCoordinator(), + LockTTL: time.Minute, + } + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeAsync, + Model: extModel, + }, + Coordination: coord, + }) + require.NoError(t, err) + + state1 := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember one"), + schema.AssistantMessage("ack1", nil), + }, + } + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{ + Messages: state1.Messages, + ToolInfos: []*schema.ToolInfo{ + {Name: "tool_one"}, + }, + }) + require.NoError(t, err) + + <-startedCh + + state2 := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember one"), + schema.AssistantMessage("ack1", nil), + schema.UserMessage("remember two"), + schema.AssistantMessage("ack2", nil), + }, + } + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{ + Messages: state2.Messages, + ToolInfos: []*schema.ToolInfo{ + {Name: "tool_one"}, + {Name: "tool_two"}, + }, + }) + require.NoError(t, err) + + close(blockCh) + + require.Eventually(t, func() bool { + topic, readErr := b.Read(ctx, &ReadRequest{FilePath: "/mem/topic.md"}) + if readErr != nil || topic == nil || topic.Content != "remember two" { + return false + } + cursor, ok, cursorErr := getCoordinatorCursor(ctx, coord.Coordinator, "/mem::session-1") + if cursorErr != nil || !ok { + return false + } + return cursor == len(state2.Messages) + }, 2*time.Second, 10*time.Millisecond) +} + +func TestMiddleware_BeforeAgent_GenInstructionRendersAndIndexInjectedOnce(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "line1\nline2\n", now) + var instructionCalls int32 + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + GenInstruction: func(ctx context.Context) (string, error) { + atomic.AddInt32(&instructionCalls, 1) + return "custom memory policy", nil + }, + // No topic selection model. + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + + _, out1, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Contains(t, out1.Instruction, "custom memory policy") + require.EqualValues(t, 1, atomic.LoadInt32(&instructionCalls)) + require.Equal(t, 1, countMemoryIndexMessages(out1.AgentInput.Messages)) + + // Same turn with already-injected index reminder should not duplicate the reminder. + _, out2, err := mw.BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: out1.Instruction, + AgentInput: &adk.AgentInput{Messages: out1.AgentInput.Messages}, + }) + require.NoError(t, err) + require.Contains(t, out2.Instruction, "custom memory policy") + require.EqualValues(t, 2, atomic.LoadInt32(&instructionCalls)) + require.Equal(t, 1, countMemoryIndexMessages(out2.AgentInput.Messages)) + + // A later business user message in the same session should not get another MEMORY.md reminder. + nextMessages := append([]*schema.Message{}, out2.AgentInput.Messages...) + nextMessages = append(nextMessages, schema.AssistantMessage("ack", nil), schema.UserMessage("next turn")) + _, out3, err := mw.BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: out2.Instruction, + AgentInput: &adk.AgentInput{Messages: nextMessages}, + }) + require.NoError(t, err) + require.Contains(t, out3.Instruction, "custom memory policy") + require.EqualValues(t, 3, atomic.LoadInt32(&instructionCalls)) + require.Equal(t, 1, countMemoryIndexMessages(out3.AgentInput.Messages)) +} + +func TestMiddleware_BeforeAgent_TopicMemoryInjectedOncePerSession(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md)\n", now) + b.put("/mem/debugging.md", "---\ndescription: debug notes\n---\nbody\n", now) + + selModel := &toolCallSelectionModel{} + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: selModel, + Read: &ReadConfig[*schema.Message]{ + Mode: ReadModeSync, + TopicSelection: &TopicSelectionConfig{ + TopK: 1, + }, + }, + }) + require.NoError(t, err) + + _, out1, err := mw.BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to debug?")}}, + }) + require.NoError(t, err) + require.EqualValues(t, 1, atomic.LoadInt32(&selModel.calls)) + require.Equal(t, 1, countMemoryIndexMessages(out1.AgentInput.Messages)) + require.Equal(t, 1, countTopicMemoryMessages(out1.AgentInput.Messages)) + + nextMessages := append([]*schema.Message{}, out1.AgentInput.Messages...) + nextMessages = append(nextMessages, schema.AssistantMessage("ack", nil), schema.UserMessage("How to debug again?")) + _, out2, err := mw.BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: out1.Instruction, + AgentInput: &adk.AgentInput{Messages: nextMessages}, + }) + require.NoError(t, err) + require.EqualValues(t, 1, atomic.LoadInt32(&selModel.calls)) + require.Equal(t, 1, countMemoryIndexMessages(out2.AgentInput.Messages)) + require.Equal(t, 1, countTopicMemoryMessages(out2.AgentInput.Messages)) +} + +func TestMiddleware_LastUserMessageSkipsSystemReminderPrefix(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &fixedModel{out: `{"selected_memories":[]}`}, + }) + require.NoError(t, err) + + last, ok := mw.(*middleware[*schema.Message]).lastUserMessage(&adk.AgentInput{ + Messages: []adk.Message{ + schema.UserMessage("real user query"), + schema.UserMessage("\nInjected by another middleware.\n"), + }, + }) + require.True(t, ok) + require.Equal(t, "real user query", last.Content) +} + +func TestMiddleware_BeforeAgent_InjectsInstructionWhenMessagesAlreadyContainMemory(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + }) + require.NoError(t, err) + + memMsg := newMemoryMessage[*schema.Message]("\n\n\nContents of /mem/preloaded.md (saved now):\npreloaded\n\n") + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi"), memMsg}}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Contains(t, out.Instruction, "# Auto memory") + require.Len(t, out.AgentInput.Messages, 3) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + require.Contains(t, out.AgentInput.Messages[1].Content, "hi") + requireTopicMemoryMessage(t, out.AgentInput.Messages[2], "preloaded") +} + +func TestMiddleware_BeforeAgent_DistributedCursorSyncIntoMessageExtra(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + coord := &CoordinationConfig[*schema.Message]{ + SessionID: "sess-cursor", + Coordinator: NewLocalCoordinator(), + LockTTL: time.Minute, + } + require.NoError(t, setCoordinatorCursor(ctx, coord.Coordinator, "/mem::sess-cursor", 5)) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Coordination: coord, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{ + schema.UserMessage("hi"), + schema.AssistantMessage("ack", nil), + }}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + requireWriteCursor(t, out.AgentInput.Messages, 5) +} + +func TestMiddleware_BeforeAgent_WriteCursorDoesNotBlockInstructionInjection(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "remembered\n", now) + + coord := &CoordinationConfig[*schema.Message]{ + SessionID: "sess-cursor", + Coordinator: NewLocalCoordinator(), + LockTTL: time.Minute, + } + require.NoError(t, setCoordinatorCursor(ctx, coord.Coordinator, "/mem::sess-cursor", 5)) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Coordination: coord, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{ + schema.AssistantMessage("ack", nil), + schema.UserMessage("next turn"), + }}, + } + + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Contains(t, out.Instruction, "# Auto memory") + require.NotContains(t, out.Instruction, "remembered") + requireMemoryIndexMessage(t, out.AgentInput.Messages[1], "remembered") + + requireWriteCursor(t, out.AgentInput.Messages, 5) +} + +func TestMiddleware_TopicSelection_ToolCallParsingAndFiltering(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md)\n", now) + b.put("/mem/debugging.md", "---\ndescription: debug notes\n---\nbody\n", now) + b.put("/mem/other.md", "---\ndescription: other\n---\nbody\n", now.Add(-time.Hour)) + + selModel := &toolCallSelectionModel{} + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: selModel, + Read: &ReadConfig[*schema.Message]{ + Mode: ReadModeSync, + TopicSelection: &TopicSelectionConfig{ + TopK: 1, + }, + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("How to debug?")}}, + } + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Len(t, out.AgentInput.Messages, 3) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + mem := out.AgentInput.Messages[1] + require.Contains(t, mem.Content, "Contents of /mem/debugging.md") + require.NotContains(t, mem.Content, "hallucinated.md") + require.Contains(t, out.AgentInput.Messages[2].Content, "How to debug?") + require.EqualValues(t, 1, atomic.LoadInt32(&selModel.calls)) +} + +func TestMiddleware_TopicSelection_AsyncProtectsMemoryMessageFromMutation(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "- [debugging.md](debugging.md)\n", now) + b.put("/mem/debugging.md", "---\ndescription: debug notes\n---\nbody\n", now) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Model: &fixedModel{out: `{"selected_memories":["debugging.md"]}`}, + Read: &ReadConfig[*schema.Message]{Mode: ReadModeAsync}, + }) + require.NoError(t, err) + + ctx2, _, err := mw.BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + }) + require.NoError(t, err) + + st := &adk.ChatModelAgentState{Messages: []adk.Message{schema.UserMessage("hi")}} + + var expected string + require.Eventually(t, func() bool { + _, next, callErr := mw.BeforeModelRewriteState(ctx2, st, nil) + require.NoError(t, callErr) + st = next + if len(st.Messages) < 2 { + return false + } + expected = st.Messages[len(st.Messages)-1].Content + return strings.Contains(expected, "") + }, 2*time.Second, 10*time.Millisecond) + + // Mutate the memory message content. + st.Messages[len(st.Messages)-1].Content = "tampered" + _, next, err := mw.BeforeModelRewriteState(ctx2, st, nil) + require.NoError(t, err) + require.Equal(t, expected, next.Messages[len(next.Messages)-1].Content) + require.NotNil(t, next.Messages[len(next.Messages)-1].Extra[memoryExtraKey]) +} + +func TestMiddleware_AfterAgent_SyncExtraction_ChinesePrompt(t *testing.T) { + require.NoError(t, adk.SetLanguage(adk.LanguageChinese)) + defer func() { + require.NoError(t, adk.SetLanguage(adk.LanguageEnglish)) + }() + + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + extModel := &extractionModel{} + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeSync, + Model: extModel, + }, + }) + require.NoError(t, err) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember chinese"), + schema.AssistantMessage("ack", nil), + }, + } + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{Messages: state.Messages}) + require.NoError(t, err) + + extModel.mu.Lock() + defer extModel.mu.Unlock() + require.NotEmpty(t, extModel.promptSeen) + require.Contains(t, extModel.promptSeen[0], "你现在扮演记忆提取子智能体") + require.Contains(t, extModel.promptSeen[0], "## 记忆目录") + require.Contains(t, extModel.promptSeen[0], "路径：/mem") +} + +func TestMiddleware_AfterAgent_RelativeMemoryDirRendersAbsolutePath(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + oldwd, err := os.Getwd() + require.NoError(t, err) + require.NoError(t, os.Chdir(tmp)) + defer func() { + _ = os.Chdir(oldwd) + }() + + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + expectedDir, err := filepath.Abs(".") + require.NoError(t, err) + + extModel := &extractionModel{} + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: ".", + MemoryBackend: NewLocalBackend(), + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeSync, + Model: extModel, + }, + }) + require.NoError(t, err) + + state := &adk.TypedChatModelAgentState[*schema.Message]{ + Messages: []adk.Message{ + schema.UserMessage("remember relative"), + schema.AssistantMessage("ack", nil), + }, + } + _, err = mw.AfterAgent(ctx, state) + require.NoError(t, err) + + extModel.mu.Lock() + require.NotEmpty(t, extModel.promptSeen) + require.Contains(t, extModel.promptSeen[0], "Path: "+expectedDir) + extModel.mu.Unlock() + + raw, err := os.ReadFile(filepath.Join(expectedDir, "topic.md")) + require.NoError(t, err) + require.Equal(t, "remember relative", string(raw)) +} + +func TestMiddleware_BeforeAgent_RelativeMemoryDirReadsResolvedDirectoryAfterCWDChange(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + oldwd, err := os.Getwd() + require.NoError(t, err) + require.NoError(t, os.Chdir(tmp)) + defer func() { + _ = os.Chdir(oldwd) + }() + + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte("persisted index\n"), 0o644)) + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: ".", + MemoryBackend: NewLocalBackend(), + }) + require.NoError(t, err) + + other := t.TempDir() + require.NoError(t, os.Chdir(other)) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}, + } + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.NotContains(t, out.Instruction, "persisted index") + require.Len(t, out.AgentInput.Messages, 2) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "persisted index") + require.Contains(t, out.AgentInput.Messages[1].Content, "hi") +} + +func TestFSBackend_ReadMissingFileReturnsContentInsteadOfError(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + + fs, err := newFSBackend(NewLocalBackend(), tmp) + require.NoError(t, err) + + content, err := fs.Read(ctx, &ReadRequest{FilePath: "missing.md"}) + require.NoError(t, err) + require.NotNil(t, content) + require.Contains(t, content.Content, "File not found:") + require.Contains(t, content.Content, filepath.Join(tmp, "missing.md")) +} + +func TestMiddleware_TopicSelection_IgnoresOutOfBoundsCandidatePaths(t *testing.T) { + ctx := context.Background() + backend := &outOfBoundsCandidateBackend{} + + mw, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: backend, + Model: &panicModel{}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + AgentInput: &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("show memories")}}, + } + _, out, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + require.Len(t, out.AgentInput.Messages, 2) + requireMemoryIndexMessage(t, out.AgentInput.Messages[0], "Contents of /mem/MEMORY.md") + require.Contains(t, out.AgentInput.Messages[1].Content, "show memories") + require.Equal(t, int32(0), atomic.LoadInt32(&backend.outsideReadCalled)) +} + +func TestMiddleware_AfterAgent_AsyncSetsPendingSnapshotWhenLockHeld(t *testing.T) { + ctx := context.Background() + b := NewInMemoryBackend() + now := time.Now() + b.put("/mem/MEMORY.md", "", now) + + extModel := &extractionModel{} + coord := &CoordinationConfig[*schema.Message]{ + SessionID: "sess-pending", + Coordinator: NewLocalCoordinator(), + LockTTL: time.Minute, + } + // Hold the lock. + coordKey := "/mem::sess-pending" + unlock, ok, err := coord.Coordinator.AcquireLock(ctx, coordKey, time.Minute) + require.NoError(t, err) + require.True(t, ok) + + mwI, err := New(ctx, &Config[*schema.Message]{ + MemoryDirectory: "/mem", + MemoryBackend: b, + Write: &WriteConfig[*schema.Message]{ + Mode: WriteModeAsync, + Model: extModel, + }, + Coordination: coord, + }) + require.NoError(t, err) + mw := mwI.(*middleware[*schema.Message]) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("remember pending"), + schema.AssistantMessage("ack", nil), + }, + } + _, err = mw.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{ + Messages: state.Messages, + ToolInfos: []*schema.ToolInfo{ + {Name: "pending_tool"}, + }, + }) + require.NoError(t, err) + + pending, err := popCoordinatorPendingSnapshot(ctx, coord.Coordinator, coordKey) + require.NoError(t, err) + require.NotNil(t, pending) + + // Release and drain manually to complete write synchronously in test. + require.NoError(t, unlock(ctx)) + unlock2, ok, err := coord.Coordinator.AcquireLock(ctx, coordKey, time.Minute) + require.NoError(t, err) + require.True(t, ok) + mw.runExtractionDrain(ctx, coordKey, unlock2, pending) + + topic, err := b.Read(ctx, &ReadRequest{FilePath: "/mem/topic.md"}) + require.NoError(t, err) + require.Equal(t, "remember pending", topic.Content) +} diff --git a/adk/middlewares/automemory/backend.go b/adk/middlewares/automemory/backend.go new file mode 100644 index 000000000..9d217351b --- /dev/null +++ b/adk/middlewares/automemory/backend.go @@ -0,0 +1,46 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +import ( + "github.com/cloudwego/eino/adk/filesystem" + ainternal "github.com/cloudwego/eino/adk/middlewares/automemory/internal" +) + +// Backend is the only filesystem storage abstraction users need to implement +// for automemory and dream. +// +// It intentionally exposes only the capabilities required by memory loading and +// consolidation: Read, GlobInfo, Write, and Edit. +// +// LocalBackend and InMemoryBackend both implement this interface. +type Backend = ainternal.Backend + +type ReadRequest = filesystem.ReadRequest +type FileContent = filesystem.FileContent +type GlobInfoRequest = filesystem.GlobInfoRequest +type FileInfo = filesystem.FileInfo +type WriteRequest = filesystem.WriteRequest +type EditRequest = filesystem.EditRequest + +func newFSBackend(backend Backend, baseDir string) (*ainternal.FSBackend, error) { + return ainternal.NewFSBackend(backend, ainternal.FSBackendConfig{ + BaseDir: baseDir, + NotFoundAsContent: true, + ErrorPrefix: "fs backend", + }) +} diff --git a/adk/middlewares/automemory/consts.go b/adk/middlewares/automemory/consts.go new file mode 100644 index 000000000..f5ee3aa43 --- /dev/null +++ b/adk/middlewares/automemory/consts.go @@ -0,0 +1,60 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +const ( + // CandidateGlobPattern matches topic files under the memory directory. + CandidateGlobPattern = "**/*.md" + + memoryIndexFileName = "MEMORY.md" + + defaultIndexMaxLines = 200 + defaultIndexMaxBytes = 4 * 1024 + + defaultCandidateLimit = 200 + defaultCandidatePreviewLine = 30 + + defaultTopicTopK = 5 + defaultTopicMaxLines = 200 + defaultTopicMaxBytes = 4 * 1024 + defaultTopicMaxTotalBytes = 16 * 1024 + + defaultMemoryWriteMaxTurns = 5 + + topicSelectionToolName = "select_memories" +) + +// ErrorStage error stage during auto memory processing +type ErrorStage string + +// OnError stage constants. These values are stable identifiers used to report +// best-effort failures through Config.OnError. +const ( + OnErrorStageTopicSelectionSync ErrorStage = "topic_selection_sync" + OnErrorStageTopicSelectionAsync ErrorStage = "topic_selection_async" + OnErrorStageRenderInstruction ErrorStage = "render_instruction" + OnErrorStageResolveSessionID ErrorStage = "resolve_session_id" + OnErrorStageMemoryWriteSync ErrorStage = "memory_write_sync" + OnErrorStageSnapshotMarshal ErrorStage = "snapshot_marshal" + OnErrorStageAcquireExtractionLock ErrorStage = "acquire_extraction_lock" + OnErrorStageStashPendingSnapshot ErrorStage = "stash_pending_snapshot" + OnErrorStageReleaseExtractionLock ErrorStage = "release_extraction_lock" + OnErrorStageDecodePendingSnapshot ErrorStage = "decode_pending_snapshot" + OnErrorStageMemoryWriteAsync ErrorStage = "memory_write_async" + OnErrorStageLoadPendingSnapshot ErrorStage = "load_pending_snapshot" + OnErrorStageSendSessionEvent ErrorStage = "send_session_event" +) diff --git a/adk/middlewares/automemory/coordinator.go b/adk/middlewares/automemory/coordinator.go new file mode 100644 index 000000000..9999aa113 --- /dev/null +++ b/adk/middlewares/automemory/coordinator.go @@ -0,0 +1,207 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "sync" + "time" + + "github.com/cloudwego/eino/adk" +) + +// Coordinator abstracts distributed coordination for async memory extraction. +// A Redis-backed implementation can map AcquireLock to SETNX + TTL, Set to SET, +// Get to GET, and GetAndDelete to GETDEL. +type Coordinator interface { + // AcquireLock tries to acquire a lock for key. When ok==true, + // it returns an unlock function that must be called exactly once. + AcquireLock(ctx context.Context, key string, ttl time.Duration) (unlock func(context.Context) error, ok bool, err error) + + // Get returns the value for key. When the key does not exist, ok is false. + Get(ctx context.Context, key string) (value []byte, ok bool, err error) + + // Set stores value for key. ttl<=0 means no expiration. + Set(ctx context.Context, key string, value []byte, ttl time.Duration) error + + // GetAndDelete returns the value for key and deletes it atomically. + // When the key does not exist, ok is false. + GetAndDelete(ctx context.Context, key string) (value []byte, ok bool, err error) +} + +type PendingSnapshot struct { + Cursor int `json:"cursor"` + Messages []byte `json:"messages"` + ToolInfos []byte `json:"tool_infos,omitempty"` +} + +type CoordinationConfig[M adk.MessageType] struct { + // SessionID is the logical session ID used to build the coordinator key. + // Optional. When empty, cross-turn coordination is disabled. + SessionID string + + // Coordinator stores cursor/pending state and coordinates async extraction locks. + // Optional. Defaults to NewLocalCoordinator(). + Coordinator Coordinator + + // LockTTL is the expiration duration for extraction locks and pending snapshots. + // Optional. Defaults to the package default lock TTL. + LockTTL time.Duration +} + +// LocalCoordinator is the default in-process coordinator used in tests and single-instance deployments. +// For distributed deployments, provide a Coordinator backed by Redis or another shared KV. +type LocalCoordinator struct { + mu sync.Mutex + locks map[string]localLock + kv map[string]localValue +} + +type localLock struct { + token string + expiry time.Time +} + +type localValue struct { + value []byte + expiry time.Time +} + +// NewLocalCoordinator returns the default in-process Coordinator implementation. +func NewLocalCoordinator() *LocalCoordinator { + return &LocalCoordinator{ + locks: map[string]localLock{}, + kv: map[string]localValue{}, + } +} + +func (c *LocalCoordinator) AcquireLock(_ context.Context, key string, ttl time.Duration) (func(context.Context) error, bool, error) { + c.mu.Lock() + defer c.mu.Unlock() + now := time.Now() + if l, ok := c.locks[key]; ok && now.Before(l.expiry) { + return nil, false, nil + } + token := randToken() + c.locks[key] = localLock{token: token, expiry: now.Add(ttl)} + return func(_ context.Context) error { + c.mu.Lock() + defer c.mu.Unlock() + l, ok := c.locks[key] + if !ok { + return nil + } + if l.token != token { + return fmt.Errorf("lock token mismatch") + } + delete(c.locks, key) + return nil + }, true, nil +} + +func (c *LocalCoordinator) Get(_ context.Context, key string) ([]byte, bool, error) { + c.mu.Lock() + defer c.mu.Unlock() + v, ok := c.kv[key] + if !ok { + return nil, false, nil + } + if !v.expiry.IsZero() && time.Now().After(v.expiry) { + delete(c.kv, key) + return nil, false, nil + } + return append([]byte(nil), v.value...), true, nil +} + +func (c *LocalCoordinator) Set(_ context.Context, key string, value []byte, ttl time.Duration) error { + c.mu.Lock() + defer c.mu.Unlock() + var expiry time.Time + if ttl > 0 { + expiry = time.Now().Add(ttl) + } + c.kv[key] = localValue{value: append([]byte(nil), value...), expiry: expiry} + return nil +} + +func (c *LocalCoordinator) GetAndDelete(_ context.Context, key string) ([]byte, bool, error) { + c.mu.Lock() + defer c.mu.Unlock() + v, ok := c.kv[key] + if !ok { + return nil, false, nil + } + delete(c.kv, key) + if !v.expiry.IsZero() && time.Now().After(v.expiry) { + return nil, false, nil + } + return append([]byte(nil), v.value...), true, nil +} + +func coordinatorCursorKey(key string) string { + return key + "::cursor" +} + +func coordinatorPendingSnapshotKey(key string) string { + return key + "::pending_snapshot" +} + +func getCoordinatorCursor(ctx context.Context, c Coordinator, key string) (int, bool, error) { + raw, ok, err := c.Get(ctx, coordinatorCursorKey(key)) + if err != nil || !ok { + return 0, ok, err + } + var cursor int + if _, err := fmt.Sscanf(string(raw), "%d", &cursor); err != nil { + return 0, false, err + } + return cursor, true, nil +} + +func setCoordinatorCursor(ctx context.Context, c Coordinator, key string, cursor int) error { + return c.Set(ctx, coordinatorCursorKey(key), []byte(fmt.Sprintf("%d", cursor)), 0) +} + +func popCoordinatorPendingSnapshot(ctx context.Context, c Coordinator, key string) (*PendingSnapshot, error) { + raw, ok, err := c.GetAndDelete(ctx, coordinatorPendingSnapshotKey(key)) + if err != nil || !ok { + return nil, err + } + var snapshot PendingSnapshot + if err := json.Unmarshal(raw, &snapshot); err != nil { + return nil, err + } + return &snapshot, nil +} + +func setCoordinatorPendingSnapshot(ctx context.Context, c Coordinator, key string, snapshot *PendingSnapshot, ttl time.Duration) error { + raw, err := json.Marshal(snapshot) + if err != nil { + return err + } + return c.Set(ctx, coordinatorPendingSnapshotKey(key), raw, ttl) +} + +func randToken() string { + var b [8]byte + _, _ = rand.Read(b[:]) + return hex.EncodeToString(b[:]) +} diff --git a/adk/middlewares/automemory/dream/api.go b/adk/middlewares/automemory/dream/api.go new file mode 100644 index 000000000..aa37c69ae --- /dev/null +++ b/adk/middlewares/automemory/dream/api.go @@ -0,0 +1,132 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/cloudwego/eino/adk" +) + +// This file holds the operations a caller invokes directly on a dream: starting a +// run, querying its status, and canceling it. The scheduled middleware is created +// with New in middleware.go. + +// Run starts a one-shot dream and returns the job id. +// +// By default (RunConfig.Sync == false) Run is asynchronous: it persists a pending +// job, starts the consolidation in a background goroutine, and returns the job id +// immediately with a nil error. Track it with GetDreamStatus and stop it with +// CancelDream using that id. Set RunConfig.Sync to block until the run reaches a +// terminal state and return its error instead. +// +// The job's lifecycle is observable through GetDreamStatus and cancelable through +// CancelDream (the in-process default store only supports same-process queries and +// requires the process to outlive the run; inject a shared store for cross-process +// visibility). +// +// Run acquires the per-memory-directory run lock so a manual dream does not write +// concurrently with a scheduled run sharing the same store; if the lock is held, Run +// returns an empty id and a nil error. The input memory directory is never modified: +// the model edits a staged working copy that is promoted to the output directory on +// success. +func Run[M adk.MessageType](ctx context.Context, cfg *RunConfig[M]) (string, error) { + cfg = cloneRunConfig(cfg) + if err := applyCoreDefaults(ctx, &cfg.BaseConfig); err != nil { + return "", err + } + m, err := newMiddleware(&cfg.BaseConfig, cfg.SessionID, nil) + if err != nil { + return "", err + } + sessionID := strings.TrimSpace(cfg.SessionID) + + var unlock func(context.Context) error + if store := m.store(); store != nil { + u, ok, lockErr := store.AcquireLock(ctx, runLockKey(m.resolvedMemoryDir), m.lockTTL()) + if lockErr != nil || !ok { + return "", lockErr + } + unlock = u + } + + job := m.newJob(sessionID, nil) + m.persistJob(ctx, job) + + if cfg.Sync { + if unlock != nil { + defer func() { _ = unlock(ctx) }() + } + if err := m.executeJob(ctx, job, sessionID, nil); err != nil { + return job.ID, err + } + return job.ID, nil + } + + // Asynchronous: detach from the request lifecycle (so the run can outlive ctx) + // while preserving its values, and hand lock ownership to the goroutine. The + // run's outcome is reported through the job record and OnError, not the return. + runCtx := withoutCancel(ctx) + go func() { + if unlock != nil { + defer func() { _ = unlock(runCtx) }() + } + if err := m.executeJob(runCtx, job, sessionID, nil); err != nil { + m.onErr(runCtx, OnErrorStageRunDream, err) + } + }() + return job.ID, nil +} + +// GetDreamStatus returns the current Job record for jobID. It returns (nil, nil) +// when no such job exists (for example after the retention TTL elapsed). +func GetDreamStatus(ctx context.Context, store KVStore, jobID string) (*Job, error) { + if store == nil { + return nil, fmt.Errorf("dream: nil store") + } + return getJob(ctx, store, jobID) +} + +// CancelDream requests cancellation of a pending or running dream job. It marks the +// job canceled in the store and signals any in-process run to abort. Canceling a job +// that has already reached a terminal state is a no-op. Cross-process runs observe +// the canceled status on their next iteration check and stop best-effort. +func CancelDream(ctx context.Context, store KVStore, jobID string) error { + if store == nil { + return fmt.Errorf("dream: nil store") + } + job, err := getJob(ctx, store, jobID) + if err != nil { + return err + } + if job == nil { + return fmt.Errorf("dream: job not found: %s", jobID) + } + if job.Status.IsTerminal() { + return nil + } + job.Status = StatusCanceled + job.EndedAt = time.Now() + if err := setJob(ctx, store, job, jobTTL); err != nil { + return err + } + signalCancel(jobID) + return nil +} diff --git a/adk/middlewares/automemory/dream/config.go b/adk/middlewares/automemory/dream/config.go new file mode 100644 index 000000000..f6e7e8755 --- /dev/null +++ b/adk/middlewares/automemory/dream/config.go @@ -0,0 +1,282 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "fmt" + "time" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/middlewares/automemory" + "github.com/cloudwego/eino/components/model" +) + +const ( + defaultMinInterval = 24 * time.Hour + defaultMinTouchedSession = 5 + defaultScanInterval = 10 * time.Minute + defaultLockTTL = time.Hour + defaultMaxConsecutiveFailures = 3 + defaultMaxIterations = 24 +) + +// errLocalKVStoreSingleProcess is surfaced through OnError when Store falls back to +// the in-process default. It is a warning, not a fatal error: dreams still run, but +// coordination (run lock, job status, scheduler touch counting) is not shared across +// instances. +var errLocalKVStoreSingleProcess = fmt.Errorf( + "dream: Store not set, using in-process KVStore; this is single-process only, so run " + + "locks and job status are not shared across instances and scheduled dreams will not " + + "trigger when the middleware is constructed per session across instances " + + "(inject a shared, durable KVStore in production)") + +// errCountGateNeedsSessionID is surfaced through OnError when the scheduled count +// gate is active (MinTouchedSession > 1) but no SessionID is configured. The touched +// set is keyed by SessionID, so without one every run records the same empty member, +// the distinct-session count is stuck at 1, and the gate never opens. It is a +// warning, not a fatal error. +var errCountGateNeedsSessionID = fmt.Errorf( + "dream: MinTouchedSession > 1 but SessionID is empty; touched sessions are " + + "counted by SessionID, so the count stays at 1 and scheduled dreams will never " + + "trigger (set a per-session SessionID, or set MinTouchedSession to 1 to gate on " + + "MinInterval alone)") + +// OnError handles non-fatal dream errors. The stage argument is one of the +// OnErrorStage* constants identifying where the failure occurred. +// Optional. Nil means ignore the error. +type OnError func(ctx context.Context, stage ErrorStage, err error) + +// HandleIterator handles the dream sub-agent event stream. The handler is +// responsible for draining the iterator (calling Next until ok==false); it may +// return an error to fail the run explicitly. Even if it does not inspect +// event.Err, dream still records the first agent error on the stream and fails the +// job accordingly, so a failed consolidation is never reported as completed. +// Optional. Nil means dream drains the iterator itself. +type HandleIterator[M adk.MessageType] func(ctx context.Context, iter *adk.AsyncIterator[*adk.TypedAgentEvent[M]]) error + +// BaseConfig holds the resources and policy shared by every dream entry point: what +// dream reads and writes, the model it runs, and how it coordinates. It carries no +// per-invocation or scheduling fields; RunConfig and MiddlewareConfig embed it and +// add those. +type BaseConfig[M adk.MessageType] struct { + // MemoryDirectory is the memory root directory. + // Required. Relative paths are resolved during init. + // The dream run reads from this directory but never modifies it: the model + // edits a working copy under StagingDirectory, which is then promoted to + // OutputDirectory. + MemoryDirectory string + + // OutputDirectory is where consolidated memory is written when a run succeeds. + // Optional. Default: MemoryDirectory (in-place consolidation). + // + // When equal to MemoryDirectory, the staged result is promoted over the source + // after the run. The source is left untouched until promotion, so a failed or + // canceled run never leaves the source half-processed. Promotion is a + // best-effort copy (not atomic across files); a warning is reported via OnError. + OutputDirectory string + + // StagingDirectory is the root under which per-run working copies are created + // (one subdirectory per job id). The model's writes are bounded to the staging + // subdirectory; the source MemoryDirectory is copied in before the run. + // Optional. Default: a directory under os.TempDir(). + StagingDirectory string + + // MemoryBackend reads and updates memory files. + // Required. + MemoryBackend automemory.Backend + + // Shell, when set, is used only to clean up the staging subdirectory after a + // successful promotion (rm -rf). + // Optional. When nil, the Shell is auto-derived from MemoryBackend if that value + // also implements filesystem.Shell (the common case for sandbox/filesystem + // backends that satisfy both interfaces in one struct), so it need not be + // configured twice. Set this field only to override that, or to supply a Shell + // when MemoryBackend is not one. When neither yields a Shell, staging directories + // are left in place under StagingDirectory. + Shell filesystem.Shell + + // Model is the model used by the internal dream agent. + // Required. + Model model.BaseModel[M] + + // MaxIterations caps the dream agent's tool-call loop. + // Consolidation reads the index, skims and reads multiple topic files, then + // writes/edits several files plus the index, so this needs headroom on large + // memory directories. + // Optional. Default: 24. + MaxIterations int + + // OnError handles non-fatal runtime errors. + // Optional. Default: nil. + OnError OnError + + // SessionStore enables session timeline lookup through `grep_session_history`. + // Optional. Default: nil. + // + // When nil, dream consolidates using only memory files plus scheduler-provided + // touch signals; no session-history search tool is exposed to the model. + // + // When set, dream exposes `grep_session_history` for the sessions included in + // the current run scope: + // - middleware-triggered runs search the touched sessions selected by the scheduler + // - manual `Run(...)` searches the provided/current session only + SessionStore adk.SessionEventStore[M] + + // Store persists job records and the per-memory-directory run lock, enabling + // GetDreamStatus/CancelDream and preventing concurrent writes to the same memory + // directory. The scheduled path additionally uses it for touched sessions and + // schedule state. + // Optional. Default: in-process store from NewLocalKVStore (single-process only; + // a warning is emitted through OnError when this default is used). + // + // See KVStore for why production deployments MUST inject a shared, durable store. + Store KVStore + + // LockTTL is the lease for the per-memory-directory run lock. + // It must comfortably exceed the longest expected dream runtime: if a run + // outlives the lease, another process may acquire the lock and write the same + // memory directory concurrently. + // Optional. Default: 1h. + LockTTL time.Duration + + // HandleIterator overrides iterator consumption. + // Optional. Default: nil. + HandleIterator HandleIterator[M] +} + +// RunConfig is the input to Run(...): a one-shot consolidation. It embeds BaseConfig +// and adds the session this run is for. +type RunConfig[M adk.MessageType] struct { + BaseConfig[M] + + // SessionID identifies the session this manual run is associated with. It scopes + // the optional grep_session_history tool. + // Optional. When empty, dream runs without session grouping. + SessionID string + + // Sync makes Run block until the consolidation reaches a terminal state and + // return its error. When false (the default), Run starts the consolidation in a + // background goroutine and returns the job id immediately with a nil error; poll + // GetDreamStatus and stop it with CancelDream using that id. + // + // Note: with the in-process default Store, an asynchronous run is only observable + // from the same process and must outlive it; inject a shared, durable Store (and + // keep the process alive) to track an async run reliably. + // Optional. Default: false (asynchronous). + Sync bool +} + +// MiddlewareConfig configures the scheduled dream middleware created by New(...). It +// embeds BaseConfig and adds the per-instance session plus the trigger knobs that +// only apply to automatic, middleware-driven runs. +type MiddlewareConfig[M adk.MessageType] struct { + BaseConfig[M] + + // SessionID identifies the session this middleware instance serves. The scheduled + // trigger counts distinct SessionIDs that have touched the memory directory, so a + // per-session SessionID is required for MinTouchedSession > 1 to be meaningful; + // without one the count stays at 1. New warns through OnError when this is + // misconfigured. + // Optional. When empty, the count gate effectively degrades to 1. + SessionID string + + // MinInterval is the minimum interval between successful runs. + // Optional. Default: 24h. + MinInterval time.Duration + + // MinTouchedSession is the minimum number of distinct sessions (by SessionID) + // that must touch the memory directory before a run. Set it to 1 to gate on + // MinInterval alone. Values > 1 require a per-session SessionID. + // Optional. Default: 5. + MinTouchedSession int + + // ScanInterval is the retry delay when the session threshold is not met. + // Optional. Default: 10m. + ScanInterval time.Duration + + // MaxConsecutiveFailures caps how many times a failing run is retried against + // the same unconsolidated window before the window is advanced to avoid + // replaying the same sessions forever. + // Optional. Default: 3. + MaxConsecutiveFailures int + + // RunInline runs triggered dreams in the `AfterAgent` call path. + // Optional. Default: false. + RunInline bool +} + +// scheduleParams carries the resolved scheduler knobs onto the middleware. It is nil +// for a Run(...)-only construction, which never triggers on a schedule. +type scheduleParams struct { + minInterval time.Duration + minTouchedSession int + scanInterval time.Duration + maxConsecutiveFailures int + runInline bool +} + +func applyCoreDefaults[M adk.MessageType](ctx context.Context, cfg *BaseConfig[M]) error { + if cfg == nil { + return fmt.Errorf("auto dream config: nil") + } + if cfg.MemoryDirectory == "" || cfg.MemoryBackend == nil || cfg.Model == nil { + return fmt.Errorf("auto dream config: invalid") + } + if cfg.LockTTL <= 0 { + cfg.LockTTL = defaultLockTTL + } + if cfg.Store == nil { + cfg.Store = NewLocalKVStore() + if cfg.OnError != nil { + cfg.OnError(ctx, OnErrorStageInit, errLocalKVStoreSingleProcess) + } + } + return nil +} + +func cloneRunConfig[M adk.MessageType](cfg *RunConfig[M]) *RunConfig[M] { + if cfg == nil { + return nil + } + cp := *cfg + return &cp +} + +func cloneMiddlewareConfig[M adk.MessageType](cfg *MiddlewareConfig[M]) *MiddlewareConfig[M] { + if cfg == nil { + return nil + } + cp := *cfg + return &cp +} + +func applyScheduleDefaults[M adk.MessageType](cfg *MiddlewareConfig[M]) { + if cfg.MinInterval <= 0 { + cfg.MinInterval = defaultMinInterval + } + if cfg.MinTouchedSession <= 0 { + cfg.MinTouchedSession = defaultMinTouchedSession + } + if cfg.ScanInterval <= 0 { + cfg.ScanInterval = defaultScanInterval + } + if cfg.MaxConsecutiveFailures <= 0 { + cfg.MaxConsecutiveFailures = defaultMaxConsecutiveFailures + } +} diff --git a/adk/middlewares/automemory/dream/consts.go b/adk/middlewares/automemory/dream/consts.go new file mode 100644 index 000000000..25f5eab34 --- /dev/null +++ b/adk/middlewares/automemory/dream/consts.go @@ -0,0 +1,71 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +// ErrorStage identifies the dream processing stage that produced a non-fatal error +// reported through Config.OnError. Dream errors are best-effort: a failure in one +// stage is surfaced through OnError and the run continues or backs off rather than +// aborting the host agent. Compare against the OnErrorStage* constants to branch on +// what went wrong. +type ErrorStage string + +// OnErrorStage constants. These values are stable identifiers used to report +// best-effort failures through Config.OnError. +const ( + // OnErrorStageInit is reported during initialization for non-fatal configuration + // warnings. The most common case is falling back to the in-process KVStore + // (errLocalKVStoreSingleProcess) when Config.Store is nil; dreams still run, but + // cross-instance coordination is unavailable. + OnErrorStageInit ErrorStage = "init" + + // OnErrorStageRecordTouch is reported when recording the current session into the + // touched-session set fails. The schedule trigger may miss this session as a + // result, but the run is otherwise unaffected. + OnErrorStageRecordTouch ErrorStage = "record_touch" + + // OnErrorStageRunDream is reported when a scheduled (middleware-triggered) dream + // run fails. The scheduler records the failure, backs off, and eventually advances + // the window after MaxConsecutiveFailures so the same sessions are not replayed + // forever. + OnErrorStageRunDream ErrorStage = "run_dream" + + // OnErrorStageSeedStaging is reported when seeding the staging working copy from + // the input memory directory fails (for example a brand-new, empty memory + // directory). The dream proceeds against whatever was seeded. + OnErrorStageSeedStaging ErrorStage = "seed_staging" + + // OnErrorStagePromote is reported around promotion of the staged result to the + // output directory. It carries errPromoteInPlaceBestEffort as a warning when the + // output directory equals the input directory (a non-atomic, in-place overwrite). + OnErrorStagePromote ErrorStage = "promote" + + // OnErrorStagePersistJob is reported when writing a Job lifecycle record to + // the store fails. The run continues, but GetDreamStatus may not reflect the + // latest status. + OnErrorStagePersistJob ErrorStage = "persist_job" + + // OnErrorStageCleanup is reported when removing the staging directory through the + // configured Shell fails after a successful promotion. The staging directory is + // left in place; it is otherwise harmless. + OnErrorStageCleanup ErrorStage = "cleanup_staging" + + // OnErrorStageToolCall is reported when a tool call inside the dream agent fails. + // The error is turned into a message and fed back to the model so it can recover + // (retry with corrected arguments or a different tool); the run is not aborted. + // It is reported for observability only. + OnErrorStageToolCall ErrorStage = "tool_call" +) diff --git a/adk/middlewares/automemory/dream/dream.go b/adk/middlewares/automemory/dream/dream.go new file mode 100644 index 000000000..23f79937e --- /dev/null +++ b/adk/middlewares/automemory/dream/dream.go @@ -0,0 +1,91 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package dream consolidates automemory-managed memory directories: it is a +// reflective "sleep" pass that an agent runs over its accumulated memory files to +// merge duplicates, drop stale entries, convert relative dates to absolute ones, and +// keep the MEMORY.md index tidy, so future sessions orient quickly. +// +// # Two ways to run +// +// Scheduled (middleware): New returns a chat-model-agent middleware that, after each +// agent run, records the session and triggers a consolidation once enough sessions +// have accumulated and a minimum interval has elapsed. Configure it with +// MiddlewareConfig. +// +// mw, err := dream.New(ctx, &dream.MiddlewareConfig[*schema.Message]{ +// BaseConfig: dream.BaseConfig[*schema.Message]{ +// MemoryDirectory: memDir, +// MemoryBackend: backend, +// Model: model, +// Store: sharedStore, // production: a shared, durable KVStore +// }, +// SessionID: sessionID, +// MinInterval: 24 * time.Hour, +// MinTouchedSession: 5, +// }) +// +// On demand: Run starts one consolidation and returns a job id. It is asynchronous +// by default — the job runs in the background and is tracked via GetDreamStatus / +// CancelDream; set RunConfig.Sync to block until it finishes. Configure it with +// RunConfig, which adds the session this run is for and the Sync flag. +// +// jobID, err := dream.Run(ctx, &dream.RunConfig[*schema.Message]{ +// BaseConfig: dream.BaseConfig[*schema.Message]{ +// MemoryDirectory: memDir, +// MemoryBackend: backend, +// Model: model, +// Store: sharedStore, +// }, +// SessionID: sessionID, +// }) +// +// # Staged, non-destructive output +// +// A run never mutates the input MemoryDirectory. The source is copied into a +// per-run working copy under StagingDirectory; the model edits only that copy; on +// success the result is promoted (a best-effort, non-atomic copy) to OutputDirectory +// (which defaults to MemoryDirectory for in-place consolidation). A failed or +// canceled run therefore leaves the source untouched. When a Shell is available +// (configured, or auto-derived from a MemoryBackend that also implements +// filesystem.Shell) the staging copy is removed afterward. +// +// # Lifecycle +// +// Both entry points create a Job record in the KVStore. GetDreamStatus polls a +// job's status (pending, running, completed, failed, canceled) and CancelDream +// aborts a pending or running one. With the in-process default store these queries +// are same-process only; inject a shared store for cross-instance visibility and for +// the run lock that prevents concurrent writes to the same directory. +// +// # Coordination store +// +// The KVStore persists job records, the run lock, and (for the scheduled path) the +// touched-session set and schedule state. The default NewLocalKVStore is +// single-process only and emits a warning through OnError; production deployments +// MUST inject a shared, durable KVStore (for example Redis-backed). See KVStore. +// +// # Source layout +// +// - api.go — Run, GetDreamStatus, CancelDream (direct operations) +// - middleware.go — New plus the scheduling/staging/promotion internals +// - config.go — BaseConfig, RunConfig, and MiddlewareConfig +// - store.go — the KVStore interface and the in-process default +// - lifecycle.go — the Job model, statuses, and cancel plumbing +// - consts.go — ErrorStage values reported through OnError +// - prompt.go — the consolidation prompt +// - session.go — the optional grep_session_history tool +package dream diff --git a/adk/middlewares/automemory/dream/dream_test.go b/adk/middlewares/automemory/dream/dream_test.go new file mode 100644 index 000000000..22c9931ac --- /dev/null +++ b/adk/middlewares/automemory/dream/dream_test.go @@ -0,0 +1,1162 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + adkfs "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/middlewares/automemory" + ainternal "github.com/cloudwego/eino/adk/middlewares/automemory/internal" + adksession "github.com/cloudwego/eino/adk/session" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type dreamModel struct { + mu sync.Mutex + prompts []string + toolNames [][]string + calls int32 +} + +func (m *dreamModel) BindTools(tools []*schema.ToolInfo) []string { + names := make([]string, 0, len(tools)) + for _, ti := range tools { + if ti != nil { + names = append(names, ti.Name) + } + } + return names +} + +func (m *dreamModel) Generate(_ context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + callCount := atomic.AddInt32(&m.calls, 1) + toolList := model.GetCommonOptions(nil, opts...).Tools + m.mu.Lock() + m.toolNames = append(m.toolNames, m.BindTools(toolList)) + for _, msg := range input { + if msg.Role == schema.User { + m.prompts = append(m.prompts, messageText(msg)) + } + } + m.mu.Unlock() + content := "" + for _, msg := range input { + if msg.Role == schema.User { + content = messageText(msg) + } + } + if callCount > 1 { + return schema.AssistantMessage("dream complete", nil), nil + } + calls := []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "read_file", Arguments: `{"file_path":"MEMORY.md"}`}}, + {ID: "2", Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file_path":"dream.md","content":"consolidated"}`}}, + {ID: "3", Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file_path":"MEMORY.md","content":"- [Dream](dream.md) - consolidated"}`}}, + } + if strings.Contains(content, "Optional session search") { + calls = append([]schema.ToolCall{{ID: "0", Function: schema.FunctionCall{Name: "grep_session_history", Arguments: `{"query":"build failure"}`}}}, calls...) + } + return schema.AssistantMessage("dream", calls), nil +} + +func (m *dreamModel) Stream(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("not implemented") +} + +func (m *dreamModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + m.mu.Lock() + m.toolNames = append(m.toolNames, m.BindTools(tools)) + m.mu.Unlock() + return m, nil +} + +func messageText(msg *schema.Message) string { + if msg == nil { + return "" + } + return msg.Content +} + +type mainAgentModel struct { + reply string +} + +func (m *mainAgentModel) Generate(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage(m.reply, nil), nil +} + +func (m *mainAgentModel) Stream(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("not implemented") +} + +func (m *mainAgentModel) WithTools([]*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +// failingDreamModel always returns an error from Generate, simulating a dream +// run that fails before it can consolidate anything. +type failingDreamModel struct { + calls int32 +} + +func (m *failingDreamModel) Generate(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m.calls, 1) + return nil, errDreamModel +} + +func (m *failingDreamModel) Stream(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errDreamModel +} + +func (m *failingDreamModel) WithTools([]*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +var errDreamModel = fmt.Errorf("dream model failure") + +// toolFailDreamModel emits an edit_file call against a nonexistent file on its first +// turn, so the tool invocation fails inside the agent loop (surfacing as event.Err) +// rather than the model call itself failing. +type toolFailDreamModel struct { + calls int32 +} + +func (m *toolFailDreamModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + if atomic.AddInt32(&m.calls, 1) == 1 { + return schema.AssistantMessage("editing", []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "edit_file", Arguments: `{"file_path":"missing.md","old_string":"a","new_string":"b"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil +} + +func (m *toolFailDreamModel) Stream(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("not implemented") +} + +func (m *toolFailDreamModel) WithTools([]*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func drainIterator(t *testing.T, iter *adk.AsyncIterator[*adk.AgentEvent]) []*adk.AgentEvent { + t.Helper() + var out []*adk.AgentEvent + for { + ev, ok := iter.Next() + if !ok { + return out + } + out = append(out, ev) + if ev != nil && ev.Err != nil { + return out + } + } +} + +type countingSessionStore struct { + adk.SessionEventStore[*schema.Message] + loadCalls int32 +} + +func (s *countingSessionStore) LoadEvents(ctx context.Context, sessionID string, req *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[*schema.Message], error) { + atomic.AddInt32(&s.loadCalls, 1) + return s.SessionEventStore.LoadEvents(ctx, sessionID, req) +} + +// newMiddlewareConfig builds a MiddlewareConfig wired for inline, low-threshold runs +// so tests trigger immediately. It keeps a separate staging dir so the input memory +// directory is never written during a run. +func newMiddlewareConfig(t *testing.T, dir string, backend automemory.Backend, m model.BaseModel[*schema.Message], store KVStore) *MiddlewareConfig[*schema.Message] { + t.Helper() + cfg := &MiddlewareConfig[*schema.Message]{ + MinInterval: time.Hour, + MinTouchedSession: 1, + ScanInterval: time.Minute, + RunInline: true, + } + cfg.MemoryDirectory = dir + cfg.StagingDirectory = t.TempDir() + cfg.MemoryBackend = backend + cfg.Model = m + cfg.Store = store + return cfg +} + +func TestBuildConsolidationPrompt_OmitsSessionSearchSectionWhenProviderMissing(t *testing.T) { + prompt := buildConsolidationPrompt("/mem", []string{"a", "b"}, false) + require.NotContains(t, prompt, "Optional session search") + require.Contains(t, prompt, "Sessions since last consolidation (2)") +} + +func TestBuildConsolidationPrompt_Chinese(t *testing.T) { + require.NoError(t, adk.SetLanguage(adk.LanguageChinese)) + defer func() { + require.NoError(t, adk.SetLanguage(adk.LanguageEnglish)) + }() + + prompt := buildConsolidationPrompt("/mem", []string{"a", "b"}, true) + require.Contains(t, prompt, "## 可选的 session 搜索") + require.Contains(t, prompt, "它只会搜索本次 dream 运行范围内包含的 session 历史") + require.Contains(t, prompt, "自上次 consolidation 以来触达过的 sessions（2）") +} + +func TestNew_DoesNotMutateConfig(t *testing.T) { + ctx := context.Background() + cfg := &MiddlewareConfig[*schema.Message]{} + cfg.MemoryDirectory = "/mem" + cfg.MemoryBackend = automemory.NewInMemoryBackend() + cfg.Model = &dreamModel{} + cfg.SessionStore = adksession.NewInMemoryStore[*schema.Message](nil) + + _, err := New(ctx, cfg) + require.NoError(t, err) + require.Empty(t, cfg.SessionID) + require.Zero(t, cfg.MinInterval) + require.Zero(t, cfg.MinTouchedSession) + require.Zero(t, cfg.ScanInterval) + require.Zero(t, cfg.LockTTL) + require.Nil(t, cfg.Store) +} + +func TestMiddleware_AfterAgent_RunInlineWithSessionStore(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte("- [Existing](existing.md) - old"), 0o644)) + store := NewLocalKVStore() + model := &dreamModel{} + eventStore := &countingSessionStore{SessionEventStore: adksession.NewInMemoryStore[*schema.Message](nil)} + err := eventStore.AppendEvents(ctx, "session-a", []*adk.SessionEvent[*schema.Message]{{ + EventID: "e1", + Kind: adk.SessionEventMessage, + Message: schema.AssistantMessage("build failure: missing dependency", nil), + }}) + require.NoError(t, err) + cfg := newMiddlewareConfig(t, tmp, automemory.NewLocalBackend(), model, store) + cfg.SessionStore = eventStore + mw, err := New(ctx, cfg) + require.NoError(t, err) + impl, ok := mw.(*middleware[*schema.Message]) + require.True(t, ok) + now := time.Now() + impl.now = func() time.Time { return now } + require.NoError(t, setScheduleState(ctx, store, impl.resolvedMemoryDir, &ScheduleState{LastConsolidatedAt: now.Add(-2 * time.Hour), NextCheckAt: now})) + require.NoError(t, store.AddToSet(ctx, touchSetKey(impl.resolvedMemoryDir), "session-a", now.Add(-30*time.Minute), time.Hour)) + + _, err = impl.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{}) + require.NoError(t, err) + + raw, err := os.ReadFile(filepath.Join(tmp, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(raw)) + require.GreaterOrEqual(t, atomic.LoadInt32(&eventStore.loadCalls), int32(1)) + model.mu.Lock() + defer model.mu.Unlock() + require.NotEmpty(t, model.prompts) + require.Contains(t, model.prompts[0], "Optional session search") +} + +func TestMiddleware_AfterAgent_FirstEligibleTouchCanTriggerImmediately(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + model := &dreamModel{} + mw, err := New(ctx, newMiddlewareConfig(t, tmp, automemory.NewLocalBackend(), model, store)) + require.NoError(t, err) + impl, ok := mw.(*middleware[*schema.Message]) + require.True(t, ok) + now := time.Now() + impl.now = func() time.Time { return now } + require.NoError(t, store.AddToSet(ctx, touchSetKey(impl.resolvedMemoryDir), "older-session", now.Add(-2*time.Minute), time.Hour)) + + _, err = impl.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{}) + require.NoError(t, err) + + raw, err := os.ReadFile(filepath.Join(tmp, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(raw)) +} + +func TestRun_ManualDreamWithoutSchedule(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + model := &dreamModel{} + + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + MemoryBackend: automemory.NewLocalBackend(), + Model: model, + }, + SessionID: "manual-session", + Sync: true, + }) + require.NoError(t, err) + require.NotEmpty(t, jobID) + + raw, err := os.ReadFile(filepath.Join(tmp, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(raw)) + model.mu.Lock() + defer model.mu.Unlock() + require.NotEmpty(t, model.prompts) + require.NotContains(t, model.prompts[0], "Sessions since last consolidation") + require.NotContains(t, model.prompts[0], "Optional session search") +} + +func TestIntegration_UserPerspective_AgentMiddlewareAutoDream(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte("- [Existing](existing.md) - old\n"), 0o644)) + + store := NewLocalKVStore() + dm := &dreamModel{} + mw, err := New(ctx, newMiddlewareConfig(t, tmp, automemory.NewLocalBackend(), dm, store)) + require.NoError(t, err) + impl, ok := mw.(*middleware[*schema.Message]) + require.True(t, ok) + now := time.Now() + impl.now = func() time.Time { return now } + require.NoError(t, store.AddToSet(ctx, touchSetKey(impl.resolvedMemoryDir), "older-session", now.Add(-3*time.Minute), time.Hour)) + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "main_agent", + Description: "main agent for dream integration test", + Model: &mainAgentModel{reply: "main answer"}, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + require.NoError(t, err) + + events := drainIterator(t, agent.Run(ctx, &adk.AgentInput{ + Messages: []adk.Message{schema.UserMessage("please help")}, + })) + require.NotEmpty(t, events) + last := events[len(events)-1] + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.Equal(t, "main answer", last.Output.MessageOutput.Message.Content) + + raw, err := os.ReadFile(filepath.Join(tmp, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(raw)) + index, err := os.ReadFile(filepath.Join(tmp, "MEMORY.md")) + require.NoError(t, err) + require.Contains(t, string(index), "dream.md") +} + +func TestIntegration_UserPerspective_RunUsesConfigSessionID(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + + model := &dreamModel{} + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + MemoryBackend: automemory.NewLocalBackend(), + Model: model, + }, + SessionID: "fallback-session", + Sync: true, + }) + require.NoError(t, err) + require.NotEmpty(t, jobID) + + raw, err := os.ReadFile(filepath.Join(tmp, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(raw)) +} + +func TestMiddleware_AfterAgent_PrunesTouchesAfterSuccess(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + mw, err := New(ctx, newMiddlewareConfig(t, tmp, automemory.NewLocalBackend(), &dreamModel{}, store)) + require.NoError(t, err) + impl := mw.(*middleware[*schema.Message]) + now := time.Now() + impl.now = func() time.Time { return now } + require.NoError(t, store.AddToSet(ctx, touchSetKey(impl.resolvedMemoryDir), "older-session", now.Add(-3*time.Minute), time.Hour)) + + _, err = impl.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{}) + require.NoError(t, err) + + // After a successful run, touches at or before LastConsolidatedAt are dropped. + remaining, err := store.ListSet(ctx, touchSetKey(impl.resolvedMemoryDir), time.Time{}) + require.NoError(t, err) + require.Empty(t, remaining) +} + +func TestMiddleware_AfterAgent_AdvancesWindowAfterMaxFailures(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + cfg := newMiddlewareConfig(t, tmp, automemory.NewLocalBackend(), &failingDreamModel{}, store) + cfg.MaxConsecutiveFailures = 2 + mw, err := New(ctx, cfg) + require.NoError(t, err) + impl := mw.(*middleware[*schema.Message]) + now := time.Now() + impl.now = func() time.Time { return now } + require.NoError(t, store.AddToSet(ctx, touchSetKey(impl.resolvedMemoryDir), "session-a", now.Add(-3*time.Minute), time.Hour)) + + // First failure: window not yet advanced, failure counted. + _, err = impl.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{}) + require.NoError(t, err) + st, err := getScheduleState(ctx, store, impl.resolvedMemoryDir) + require.NoError(t, err) + require.NotNil(t, st) + require.Equal(t, 1, st.ConsecutiveFailures) + require.True(t, st.LastConsolidatedAt.IsZero()) + + // Allow the next check to fire and retry. + require.NoError(t, setScheduleState(ctx, store, impl.resolvedMemoryDir, &ScheduleState{ConsecutiveFailures: 1, NextCheckAt: now})) + + // Second failure reaches MaxConsecutiveFailures: window advances and counter resets. + _, err = impl.AfterAgent(ctx, &adk.TypedChatModelAgentState[*schema.Message]{}) + require.NoError(t, err) + st, err = getScheduleState(ctx, store, impl.resolvedMemoryDir) + require.NoError(t, err) + require.NotNil(t, st) + require.Equal(t, 0, st.ConsecutiveFailures) + require.True(t, st.LastConsolidatedAt.Equal(now)) +} + +func TestRun_InputDirUnchangedAndOutputPopulated(t *testing.T) { + ctx := context.Background() + input := t.TempDir() + output := t.TempDir() + staging := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(input, "MEMORY.md"), []byte("- [Existing](existing.md) - old\n"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(input, "existing.md"), []byte("original"), 0o644)) + + store := NewLocalKVStore() + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: input, + OutputDirectory: output, + StagingDirectory: staging, + MemoryBackend: automemory.NewLocalBackend(), + Model: &dreamModel{}, + Store: store, + }, + SessionID: "s1", + Sync: true, + }) + require.NoError(t, err) + require.NotEmpty(t, jobID) + + // Input directory is untouched. + existing, err := os.ReadFile(filepath.Join(input, "existing.md")) + require.NoError(t, err) + require.Equal(t, "original", string(existing)) + _, err = os.Stat(filepath.Join(input, "dream.md")) + require.True(t, os.IsNotExist(err), "input dir must not gain consolidated files") + + // Output directory has the consolidated result plus the seeded copy. + consolidated, err := os.ReadFile(filepath.Join(output, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(consolidated)) + seeded, err := os.ReadFile(filepath.Join(output, "existing.md")) + require.NoError(t, err) + require.Equal(t, "original", string(seeded)) + + // Job reached completed and is queryable. + job, err := GetDreamStatus(ctx, store, jobID) + require.NoError(t, err) + require.NotNil(t, job) + require.Equal(t, StatusCompleted, job.Status) + require.Equal(t, input, job.InputDir) + require.Equal(t, output, job.OutputDir) +} + +func TestRun_OnlyCopiesMarkdown(t *testing.T) { + ctx := context.Background() + input := t.TempDir() + output := t.TempDir() + staging := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(input, "MEMORY.md"), []byte(""), 0o644)) + // Non-markdown files mixed into the memory directory must not be copied. + require.NoError(t, os.WriteFile(filepath.Join(input, "data.json"), []byte(`{"k":"v"}`), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(input, "notes.txt"), []byte("scratch"), 0o644)) + + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: input, + OutputDirectory: output, + StagingDirectory: staging, + MemoryBackend: automemory.NewLocalBackend(), + Model: &dreamModel{}, + Store: NewLocalKVStore(), + }, + SessionID: "s1", + Sync: true, + }) + require.NoError(t, err) + + // Markdown is consolidated into the output. + consolidated, err := os.ReadFile(filepath.Join(output, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(consolidated)) + + // Non-markdown files are neither staged nor promoted. + for _, name := range []string{"data.json", "notes.txt"} { + _, statErr := os.Stat(filepath.Join(staging, jobID, name)) + require.True(t, os.IsNotExist(statErr), "staging must not contain %s", name) + _, statErr = os.Stat(filepath.Join(output, name)) + require.True(t, os.IsNotExist(statErr), "output must not contain %s", name) + } + // They remain untouched in the input directory. + raw, err := os.ReadFile(filepath.Join(input, "data.json")) + require.NoError(t, err) + require.Equal(t, `{"k":"v"}`, string(raw)) +} + +func TestRun_FailedJobStatus(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + StagingDirectory: t.TempDir(), + MemoryBackend: automemory.NewLocalBackend(), + Model: &failingDreamModel{}, + Store: store, + }, + SessionID: "s1", + Sync: true, + }) + require.Error(t, err) + require.NotEmpty(t, jobID) + + job, err := GetDreamStatus(ctx, store, jobID) + require.NoError(t, err) + require.NotNil(t, job) + require.Equal(t, StatusFailed, job.Status) + require.NotEmpty(t, job.ErrMsg) +} + +func TestRun_ToolErrorRecoveredAndReported(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + + var mu sync.Mutex + var toolErrs int + model := &toolFailDreamModel{} + + // The model's first turn issues an edit_file that fails; recovery feeds the error + // back and the model finishes on its next turn. The run completes, and the tool + // failure is reported (once) through OnError. + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + StagingDirectory: t.TempDir(), + MemoryBackend: automemory.NewLocalBackend(), + Model: model, + Store: store, + OnError: func(_ context.Context, stage ErrorStage, _ error) { + if stage == OnErrorStageToolCall { + mu.Lock() + toolErrs++ + mu.Unlock() + } + }, + }, + SessionID: "s1", + Sync: true, + }) + require.NoError(t, err, "a recoverable tool error must not fail the run") + require.NotEmpty(t, jobID) + + job, err := GetDreamStatus(ctx, store, jobID) + require.NoError(t, err) + require.NotNil(t, job) + require.Equal(t, StatusCompleted, job.Status) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, 1, toolErrs, "the tool failure should be reported exactly once via OnError") + require.GreaterOrEqual(t, atomic.LoadInt32(&model.calls), int32(2), "model should be re-invoked after the tool error") +} + +func TestRun_ModelErrorFailsJobEvenWhenHandleIteratorIgnoresErr(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + + // A HandleIterator that drains events but never inspects event.Err. A fatal + // (non-tool) error — here a model-call failure — must still fail the job. + handler := func(_ context.Context, iter *adk.AsyncIterator[*adk.AgentEvent]) error { + for { + _, ok := iter.Next() + if !ok { + return nil + } + } + } + + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + StagingDirectory: t.TempDir(), + MemoryBackend: automemory.NewLocalBackend(), + Model: &failingDreamModel{}, + Store: store, + HandleIterator: handler, + }, + SessionID: "s1", + Sync: true, + }) + require.Error(t, err, "a fatal model error must surface as a run error") + require.NotEmpty(t, jobID) + + job, err := GetDreamStatus(ctx, store, jobID) + require.NoError(t, err) + require.NotNil(t, job) + require.Equal(t, StatusFailed, job.Status) + require.NotEmpty(t, job.ErrMsg) +} + +func TestRun_AsyncByDefaultReturnsImmediatelyAndCompletes(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + output := t.TempDir() + store := NewLocalKVStore() + + // Default (Sync == false): Run returns the job id immediately with a nil error, + // before the consolidation finishes. + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + OutputDirectory: output, + StagingDirectory: t.TempDir(), + MemoryBackend: automemory.NewLocalBackend(), + Model: &dreamModel{}, + Store: store, + }, + SessionID: "s1", + }) + require.NoError(t, err) + require.NotEmpty(t, jobID) + + // The job is observable and eventually reaches completed via polling. + waitForJobStatus(t, ctx, store, jobID, StatusCompleted) + + consolidated, err := os.ReadFile(filepath.Join(output, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(consolidated)) +} + +func TestCancelDream_TerminatesRunningJobStatus(t *testing.T) { + ctx := context.Background() + store := NewLocalKVStore() + job := &Job{ID: "drm_test", Status: StatusRunning, CreatedAt: time.Now()} + require.NoError(t, setJob(ctx, store, job, time.Hour)) + + require.NoError(t, CancelDream(ctx, store, "drm_test")) + + got, err := GetDreamStatus(ctx, store, "drm_test") + require.NoError(t, err) + require.NotNil(t, got) + require.Equal(t, StatusCanceled, got.Status) + require.False(t, got.EndedAt.IsZero()) + + // Canceling a terminal job is a no-op. + require.NoError(t, CancelDream(ctx, store, "drm_test")) +} + +func TestLocalKVStore_SetOps(t *testing.T) { + ctx := context.Background() + store := NewLocalKVStore() + base := time.Now() + require.NoError(t, store.AddToSet(ctx, "k", "a", base.Add(-10*time.Minute), time.Hour)) + require.NoError(t, store.AddToSet(ctx, "k", "b", base.Add(-1*time.Minute), time.Hour)) + + all, err := store.ListSet(ctx, "k", time.Time{}) + require.NoError(t, err) + require.ElementsMatch(t, []string{"a", "b"}, all) + + recent, err := store.ListSet(ctx, "k", base.Add(-5*time.Minute)) + require.NoError(t, err) + require.Equal(t, []string{"b"}, recent) + + require.NoError(t, store.PruneSet(ctx, "k", base.Add(-5*time.Minute))) + left, err := store.ListSet(ctx, "k", time.Time{}) + require.NoError(t, err) + require.Equal(t, []string{"b"}, left) +} + +func TestNew_WarnsOnLocalStoreDefault(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + var mu sync.Mutex + var warnings []error + cfg := &MiddlewareConfig[*schema.Message]{} + cfg.MemoryDirectory = tmp + cfg.MemoryBackend = automemory.NewLocalBackend() + cfg.Model = &dreamModel{} + cfg.SessionID = "s1" // isolate the store warning from the count-gate warning + cfg.OnError = func(_ context.Context, stage ErrorStage, err error) { + if stage == OnErrorStageInit { + mu.Lock() + warnings = append(warnings, err) + mu.Unlock() + } + } + _, err := New(ctx, cfg) + require.NoError(t, err) + mu.Lock() + defer mu.Unlock() + require.Len(t, warnings, 1) + require.ErrorIs(t, warnings[0], errLocalKVStoreSingleProcess) +} + +func TestNew_WarnsWhenCountGateLacksSessionID(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + var mu sync.Mutex + var gotCountGate bool + cfg := &MiddlewareConfig[*schema.Message]{MinTouchedSession: 5} + cfg.MemoryDirectory = tmp + cfg.MemoryBackend = automemory.NewLocalBackend() + cfg.Model = &dreamModel{} + cfg.Store = NewLocalKVStore() // avoid the unrelated store warning + // SessionID intentionally empty. + cfg.OnError = func(_ context.Context, _ ErrorStage, err error) { + if err == errCountGateNeedsSessionID { + mu.Lock() + gotCountGate = true + mu.Unlock() + } + } + _, err := New(ctx, cfg) + require.NoError(t, err) + mu.Lock() + defer mu.Unlock() + require.True(t, gotCountGate, "expected count-gate warning when MinTouchedSession>1 and SessionID empty") +} + +func TestNew_NoCountGateWarningWhenSatisfiable(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + var mu sync.Mutex + var gotCountGate bool + record := func(_ context.Context, _ ErrorStage, err error) { + if err == errCountGateNeedsSessionID { + mu.Lock() + gotCountGate = true + mu.Unlock() + } + } + + // Case 1: SessionID set with the count gate active -> no warning. + cfg := &MiddlewareConfig[*schema.Message]{MinTouchedSession: 5} + cfg.MemoryDirectory = tmp + cfg.MemoryBackend = automemory.NewLocalBackend() + cfg.Model = &dreamModel{} + cfg.Store = NewLocalKVStore() + cfg.SessionID = "s1" + cfg.OnError = record + _, err := New(ctx, cfg) + require.NoError(t, err) + + // Case 2: count gate disabled (MinTouchedSession=1), no SessionID -> no warning. + cfg2 := &MiddlewareConfig[*schema.Message]{MinTouchedSession: 1} + cfg2.MemoryDirectory = tmp + cfg2.MemoryBackend = automemory.NewLocalBackend() + cfg2.Model = &dreamModel{} + cfg2.Store = NewLocalKVStore() + cfg2.OnError = record + _, err = New(ctx, cfg2) + require.NoError(t, err) + + mu.Lock() + defer mu.Unlock() + require.False(t, gotCountGate, "count-gate warning should not fire when the gate is satisfiable or disabled") +} + +// recordingShell captures the commands it is asked to execute so cleanup behavior +// can be asserted without touching the real filesystem. +type recordingShell struct { + mu sync.Mutex + commands []string +} + +func (s *recordingShell) Execute(_ context.Context, req *adkfs.ExecuteRequest) (*adkfs.ExecuteResponse, error) { + s.mu.Lock() + s.commands = append(s.commands, req.Command) + s.mu.Unlock() + code := 0 + return &adkfs.ExecuteResponse{ExitCode: &code}, nil +} + +func (s *recordingShell) snapshot() []string { + s.mu.Lock() + defer s.mu.Unlock() + return append([]string(nil), s.commands...) +} + +func TestRun_ShellCleansUpStagingAfterPromote(t *testing.T) { + ctx := context.Background() + input := t.TempDir() + output := t.TempDir() + staging := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(input, "MEMORY.md"), []byte(""), 0o644)) + shell := &recordingShell{} + + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: input, + OutputDirectory: output, + StagingDirectory: staging, + MemoryBackend: automemory.NewLocalBackend(), + Model: &dreamModel{}, + Shell: shell, + Store: NewLocalKVStore(), + }, + SessionID: "s1", + Sync: true, + }) + require.NoError(t, err) + + cmds := shell.snapshot() + require.Len(t, cmds, 1) + require.Contains(t, cmds[0], "rm -rf") + require.Contains(t, cmds[0], filepath.Join(staging, jobID)) +} + +// backendWithShell embeds a Backend and a Shell in one struct, mirroring a +// sandbox/filesystem implementation that satisfies both interfaces. It lets the test +// verify the Shell is auto-derived from MemoryBackend without configuring it twice. +type backendWithShell struct { + automemory.Backend + shell *recordingShell +} + +func (b *backendWithShell) Execute(ctx context.Context, req *adkfs.ExecuteRequest) (*adkfs.ExecuteResponse, error) { + return b.shell.Execute(ctx, req) +} + +func TestRun_ShellAutoDerivedFromBackend(t *testing.T) { + ctx := context.Background() + input := t.TempDir() + output := t.TempDir() + staging := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(input, "MEMORY.md"), []byte(""), 0o644)) + shell := &recordingShell{} + backend := &backendWithShell{Backend: automemory.NewLocalBackend(), shell: shell} + + // Config.Shell is intentionally left nil; the Shell must be derived from the + // MemoryBackend because it also implements filesystem.Shell. + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: input, + OutputDirectory: output, + StagingDirectory: staging, + MemoryBackend: backend, + Model: &dreamModel{}, + Store: NewLocalKVStore(), + }, + SessionID: "s1", + Sync: true, + }) + require.NoError(t, err) + + cmds := shell.snapshot() + require.Len(t, cmds, 1) + require.Contains(t, cmds[0], "rm -rf") + require.Contains(t, cmds[0], filepath.Join(staging, jobID)) +} + +func TestRun_NoShellLeavesStagingInPlace(t *testing.T) { + ctx := context.Background() + input := t.TempDir() + output := t.TempDir() + staging := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(input, "MEMORY.md"), []byte(""), 0o644)) + + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: input, + OutputDirectory: output, + StagingDirectory: staging, + MemoryBackend: automemory.NewLocalBackend(), + Model: &dreamModel{}, + Store: NewLocalKVStore(), + }, + SessionID: "s1", + Sync: true, + }) + require.NoError(t, err) + + // Without a Shell, the staging directory and its consolidated output remain. + staged, err := os.ReadFile(filepath.Join(staging, jobID, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(staged)) +} + +func TestRun_InPlacePromoteWarns(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + + var mu sync.Mutex + var promoteWarns []error + _, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, // OutputDirectory defaults to MemoryDirectory -> in place + StagingDirectory: t.TempDir(), + MemoryBackend: automemory.NewLocalBackend(), + Model: &dreamModel{}, + OnError: func(_ context.Context, stage ErrorStage, err error) { + if stage == OnErrorStagePromote { + mu.Lock() + promoteWarns = append(promoteWarns, err) + mu.Unlock() + } + }, + Store: NewLocalKVStore(), + }, + SessionID: "s1", + Sync: true, + }) + require.NoError(t, err) + + // The consolidated result lands in the source directory. + raw, err := os.ReadFile(filepath.Join(tmp, "dream.md")) + require.NoError(t, err) + require.Equal(t, "consolidated", string(raw)) + + mu.Lock() + defer mu.Unlock() + require.Len(t, promoteWarns, 1) + require.ErrorIs(t, promoteWarns[0], errPromoteInPlaceBestEffort) +} + +func TestRun_ReturnsEmptyWhenRunLockHeld(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + + resolved, err := ainternal.ResolveMemoryDir(tmp) + require.NoError(t, err) + // Pre-acquire the run lock another runner would hold. + _, ok, err := store.AcquireLock(ctx, runLockKey(resolved), time.Hour) + require.NoError(t, err) + require.True(t, ok) + + jobID, err := Run(ctx, &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + StagingDirectory: t.TempDir(), + MemoryBackend: automemory.NewLocalBackend(), + Model: &dreamModel{}, + Store: store, + }, + SessionID: "s1", + }) + require.NoError(t, err) + require.Empty(t, jobID, "Run must not run while the run lock is held") + + // The dream did not execute, so no consolidated file was produced. + _, statErr := os.Stat(filepath.Join(tmp, "dream.md")) + require.True(t, os.IsNotExist(statErr)) +} + +func TestGetDreamStatus_UnknownJobReturnsNil(t *testing.T) { + ctx := context.Background() + store := NewLocalKVStore() + job, err := GetDreamStatus(ctx, store, "drm_missing") + require.NoError(t, err) + require.Nil(t, job) +} + +func TestCancelDream_UnknownJobErrors(t *testing.T) { + ctx := context.Background() + store := NewLocalKVStore() + err := CancelDream(ctx, store, "drm_missing") + require.Error(t, err) +} + +func TestLifecycle_NilStoreGuards(t *testing.T) { + ctx := context.Background() + _, err := GetDreamStatus(ctx, nil, "drm_x") + require.Error(t, err) + require.Error(t, CancelDream(ctx, nil, "drm_x")) +} + +func TestCancelDream_AbortsRunningJobBeforeCompletion(t *testing.T) { + ctx := context.Background() + tmp := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(tmp, "MEMORY.md"), []byte(""), 0o644)) + store := NewLocalKVStore() + + // gateModel blocks on the first generation until released, giving the test time + // to cancel an in-flight run; the cancel is observed via the run context. + gm := &gateModel{released: make(chan struct{})} + + cfg := &RunConfig[*schema.Message]{ + BaseConfig: BaseConfig[*schema.Message]{ + MemoryDirectory: tmp, + StagingDirectory: t.TempDir(), + MemoryBackend: automemory.NewLocalBackend(), + Model: gm, + Store: store, + }, + SessionID: "s1", + } + + // Run is asynchronous by default: it returns immediately while the run blocks in + // gateModel. + jobID, err := Run(ctx, cfg) + require.NoError(t, err) + require.NotEmpty(t, jobID) + + // Wait until the run is in progress, then cancel it and let the model proceed. + running := waitForRunningJob(t, ctx, store) + require.Equal(t, jobID, running) + require.NoError(t, CancelDream(ctx, store, jobID)) + close(gm.released) + + // The job settles into canceled. + waitForJobStatus(t, ctx, store, jobID, StatusCanceled) + + // A canceled run must not promote a consolidated result. + _, statErr := os.Stat(filepath.Join(tmp, "dream.md")) + require.True(t, os.IsNotExist(statErr)) +} + +// gateModel blocks the first Generate call until released is closed, then drives a +// normal consolidation. It lets a test observe and cancel a running dream. +type gateModel struct { + released chan struct{} + calls int32 +} + +func (m *gateModel) Generate(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + if atomic.AddInt32(&m.calls, 1) == 1 { + select { + case <-m.released: + case <-ctx.Done(): + return nil, ctx.Err() + } + return schema.AssistantMessage("dream", []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file_path":"dream.md","content":"consolidated"}`}}, + }), nil + } + return schema.AssistantMessage("dream complete", nil), nil +} + +func (m *gateModel) Stream(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("not implemented") +} + +func (m *gateModel) WithTools([]*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func waitForRunningJob(t *testing.T, ctx context.Context, store KVStore) string { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if id := scanForRunningJob(ctx, store); id != "" { + return id + } + time.Sleep(5 * time.Millisecond) + } + t.Fatal("no running dream job appeared") + return "" +} + +// waitForJobStatus polls until the job reaches want or the deadline elapses. +func waitForJobStatus(t *testing.T, ctx context.Context, store KVStore, jobID string, want Status) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + job, err := GetDreamStatus(ctx, store, jobID) + require.NoError(t, err) + if job != nil && job.Status == want { + return + } + time.Sleep(5 * time.Millisecond) + } + job, _ := GetDreamStatus(ctx, store, jobID) + var got Status + if job != nil { + got = job.Status + } + t.Fatalf("job %s did not reach status %q (last: %q)", jobID, want, got) +} + +// scanForRunningJob looks up the running job by probing the local store's job keys. +func scanForRunningJob(ctx context.Context, store KVStore) string { + ls, ok := store.(*localKVStore) + if !ok { + return "" + } + ls.mu.Lock() + defer ls.mu.Unlock() + for key := range ls.kv { + if !strings.HasPrefix(key, "dream::job::") { + continue + } + id := strings.TrimPrefix(key, "dream::job::") + if job, _ := getJobLocked(ls, id); job != nil && job.Status == StatusRunning { + return id + } + } + return "" +} + +// getJobLocked reads a job from an already-locked localKVStore. +func getJobLocked(ls *localKVStore, jobID string) (*Job, error) { + v, ok := ls.kv[jobKey(jobID)] + if !ok { + return nil, nil + } + var job Job + if err := json.Unmarshal(v.value, &job); err != nil { + return nil, err + } + return &job, nil +} diff --git a/adk/middlewares/automemory/dream/lifecycle.go b/adk/middlewares/automemory/dream/lifecycle.go new file mode 100644 index 000000000..a255ae569 --- /dev/null +++ b/adk/middlewares/automemory/dream/lifecycle.go @@ -0,0 +1,182 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "fmt" + "sync" + "time" +) + +// This file defines the dream job data model and the in-process cancel plumbing the +// operations in api.go act on. + +// Status is the lifecycle status of a dream job. +type Status string + +const ( + // StatusPending means the job was created and queued. + StatusPending Status = "pending" + // StatusRunning means the consolidation pipeline is processing. + StatusRunning Status = "running" + // StatusCompleted means the job finished successfully and the output + // directory holds the consolidated memory. + StatusCompleted Status = "completed" + // StatusFailed means the job terminated with an error. + StatusFailed Status = "failed" + // StatusCanceled means the job was canceled before completing. + StatusCanceled Status = "canceled" +) + +// IsTerminal reports whether the status is a terminal state. +func (s Status) IsTerminal() bool { + switch s { + case StatusCompleted, StatusFailed, StatusCanceled: + return true + default: + return false + } +} + +// Job is the observable record of one dream run. It is persisted to the +// KVStore so callers can poll GetDreamStatus and request CancelDream. +type Job struct { + // ID is the unique job identifier returned by Run and by middleware-triggered runs. + ID string `json:"id"` + + // Status is the current lifecycle status. + Status Status `json:"status"` + + // InputDir is the source memory directory the run reads from. It is never modified. + InputDir string `json:"input_dir"` + + // OutputDir is where consolidated memory is written on success. + OutputDir string `json:"output_dir"` + + // StagingDir is the working copy the model edits during the run. + StagingDir string `json:"staging_dir"` + + // SessionIDs are the sessions in scope for this run. + SessionIDs []string `json:"session_ids,omitempty"` + + // CreatedAt is when the job record was created. + CreatedAt time.Time `json:"created_at"` + + // EndedAt is when the job reached a terminal state. Zero while non-terminal. + EndedAt time.Time `json:"ended_at,omitempty"` + + // ErrMsg is the failure reason when Status is failed. Empty otherwise. + ErrMsg string `json:"err_msg,omitempty"` +} + +// jobTTL is how long terminal job records are retained for status queries. +const jobTTL = 24 * time.Hour + +// cancelRegistry tracks the cancel handle for jobs running in this process so +// CancelDream can abort an in-flight run. Cross-process cancellation is +// best-effort: a running node also polls its job record between agent iterations. +var cancelRegistry sync.Map // jobID -> *cancelHandle + +var errDreamCanceled = fmt.Errorf("dream: canceled") + +// cancelHandle pairs a context.CancelFunc with the cause it was canceled for. +// It stands in for Go 1.21's context.WithCancelCause/context.Cause, which are not +// available on the Go 1.18 baseline this module supports. +type cancelHandle struct { + cancel context.CancelFunc + + mu sync.Mutex + cause error +} + +// trigger records the cause (first writer wins) and cancels the context. +func (h *cancelHandle) trigger(cause error) { + h.mu.Lock() + if h.cause == nil { + h.cause = cause + } + h.mu.Unlock() + h.cancel() +} + +// Cause returns the recorded cancel cause, or nil if none was set. +func (h *cancelHandle) Cause() error { + h.mu.Lock() + defer h.mu.Unlock() + return h.cause +} + +func registerCancel(jobID string, h *cancelHandle) { + cancelRegistry.Store(jobID, h) +} + +func unregisterCancel(jobID string) { + cancelRegistry.Delete(jobID) +} + +func signalCancel(jobID string) { + if v, ok := cancelRegistry.Load(jobID); ok { + if h, ok := v.(*cancelHandle); ok { + h.trigger(errDreamCanceled) + } + } +} + +// jobCancelCause returns the cancel cause recorded for jobID if its run is still +// tracked in this process; otherwise it falls back to ctx.Err(). It lets a canceled +// run distinguish a dream cancellation (errDreamCanceled) from an external parent +// cancellation (context.Canceled) without context.Cause. +func jobCancelCause(jobID string, ctx context.Context) error { + if v, ok := cancelRegistry.Load(jobID); ok { + if h, ok := v.(*cancelHandle); ok { + if cause := h.Cause(); cause != nil { + return cause + } + } + } + return ctx.Err() +} + +// detachedContext carries a parent's values but is never canceled and has no +// deadline. It replaces context.WithoutCancel (Go 1.21) on the Go 1.18 baseline, +// letting a background dream run outlive the request context while still reading its +// values (loggers, trace spans). +type detachedContext struct { + parent context.Context +} + +func withoutCancel(parent context.Context) context.Context { + return detachedContext{parent: parent} +} + +func (detachedContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (detachedContext) Done() <-chan struct{} { return nil } +func (detachedContext) Err() error { return nil } +func (c detachedContext) Value(key any) any { + if c.parent == nil { + return nil + } + return c.parent.Value(key) +} + +// newJobID returns a unique-ish job id. It does not rely on time/random sources +// being deterministic; the random token plus the supplied seed make collisions +// negligible in practice. +func newJobID(seed string) string { + return "drm_" + dirHash(seed+randToken()) +} diff --git a/adk/middlewares/automemory/dream/middleware.go b/adk/middlewares/automemory/dream/middleware.go new file mode 100644 index 000000000..62c99f026 --- /dev/null +++ b/adk/middlewares/automemory/dream/middleware.go @@ -0,0 +1,620 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/cloudwego/eino/adk" + adkfs "github.com/cloudwego/eino/adk/filesystem" + ainternal "github.com/cloudwego/eino/adk/middlewares/automemory/internal" + fsmw "github.com/cloudwego/eino/adk/middlewares/filesystem" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// copyGlobPattern matches the files dream stages and promotes. It is limited to +// markdown so that, when a memory directory mixes memory files with unrelated files, +// only the memory files are copied into the working copy and back out — matching the +// memory scope used elsewhere (automemory.CandidateGlobPattern). +const copyGlobPattern = "**/*.md" + +// cancelPollInterval is how many agent events pass between cross-process cancel +// checks (re-reading the job status from the store) during a drained run. +const cancelPollInterval = 8 + +// errPromoteInPlaceBestEffort warns that an in-place promotion (OutputDirectory == +// MemoryDirectory) is a non-atomic, best-effort copy. The source is untouched until +// promotion, but a crash mid-copy can leave the live directory partially updated. +var errPromoteInPlaceBestEffort = fmt.Errorf( + "dream: promoting in place (output dir == memory dir) via non-atomic copy; " + + "a crash mid-promotion can leave the directory partially updated") + +// New creates middleware that triggers dream automatically after agent runs. +// +// When MiddlewareConfig.Store is nil, an in-process store is used and a warning is +// reported through OnError; see KVStore for why production deployments must inject a +// shared, durable store. +func New[M adk.MessageType](ctx context.Context, cfg *MiddlewareConfig[M]) (adk.TypedChatModelAgentMiddleware[M], error) { + if cfg == nil { + return nil, fmt.Errorf("auto dream config: nil") + } + cfg = cloneMiddlewareConfig(cfg) + if err := applyCoreDefaults(ctx, &cfg.BaseConfig); err != nil { + return nil, err + } + applyScheduleDefaults(cfg) + // The touched set is keyed by SessionID. With the count gate active but no + // SessionID, the distinct-session count is pinned at 1 and dreams never trigger; + // warn so the misconfiguration is visible rather than silent. + if cfg.MinTouchedSession > 1 && strings.TrimSpace(cfg.SessionID) == "" && cfg.OnError != nil { + cfg.OnError(ctx, OnErrorStageInit, errCountGateNeedsSessionID) + } + return newMiddleware(&cfg.BaseConfig, cfg.SessionID, &scheduleParams{ + minInterval: cfg.MinInterval, + minTouchedSession: cfg.MinTouchedSession, + scanInterval: cfg.ScanInterval, + maxConsecutiveFailures: cfg.MaxConsecutiveFailures, + runInline: cfg.RunInline, + }) +} + +type middleware[M adk.MessageType] struct { + adk.TypedBaseChatModelAgentMiddleware[M] + + cfg *BaseConfig[M] + sched *scheduleParams + sessionID string + resolvedMemoryDir string + resolvedOutputDir string + stagingRoot string + inputFS *ainternal.FSBackend + outputFS *ainternal.FSBackend + shell adkfs.Shell + sessionSearchTool tool.BaseTool + now func() time.Time +} + +func newMiddleware[M adk.MessageType](cfg *BaseConfig[M], sessionID string, sched *scheduleParams) (*middleware[M], error) { + resolvedMemoryDir, err := ainternal.ResolveMemoryDir(cfg.MemoryDirectory) + if err != nil { + return nil, fmt.Errorf("auto dream config: resolve memory dir: %w", err) + } + outputDir := cfg.OutputDirectory + if strings.TrimSpace(outputDir) == "" { + outputDir = cfg.MemoryDirectory + } + resolvedOutputDir, err := ainternal.ResolveMemoryDir(outputDir) + if err != nil { + return nil, fmt.Errorf("auto dream config: resolve output dir: %w", err) + } + stagingRoot := cfg.StagingDirectory + if strings.TrimSpace(stagingRoot) == "" { + stagingRoot = filepath.Join(os.TempDir(), "eino-dream") + } + resolvedStagingRoot, err := ainternal.ResolveMemoryDir(stagingRoot) + if err != nil { + return nil, fmt.Errorf("auto dream config: resolve staging dir: %w", err) + } + + inputFS, err := ainternal.NewFSBackend(cfg.MemoryBackend, ainternal.FSBackendConfig{ + BaseDir: resolvedMemoryDir, + AllowLs: true, + NotFoundAsContent: true, + ErrorPrefix: "dream input backend", + }) + if err != nil { + return nil, err + } + outputFS, err := ainternal.NewFSBackend(cfg.MemoryBackend, ainternal.FSBackendConfig{ + BaseDir: resolvedOutputDir, + AllowLs: true, + NotFoundAsContent: true, + ErrorPrefix: "dream output backend", + }) + if err != nil { + return nil, err + } + + var sessionSearchTool tool.BaseTool + if cfg.SessionStore != nil { + sessionSearchTool, err = newSessionHistoryGrepTool(cfg.SessionStore) + if err != nil { + return nil, err + } + } + m := &middleware[M]{ + cfg: cfg, + sched: sched, + sessionID: strings.TrimSpace(sessionID), + resolvedMemoryDir: resolvedMemoryDir, + resolvedOutputDir: resolvedOutputDir, + stagingRoot: resolvedStagingRoot, + inputFS: inputFS, + outputFS: outputFS, + shell: resolveShell(cfg), + sessionSearchTool: sessionSearchTool, + now: time.Now, + } + return m, nil +} + +// resolveShell returns the Shell used for staging cleanup. An explicit Config.Shell +// wins; otherwise the MemoryBackend is reused when it also implements filesystem.Shell, +// so a backend struct that satisfies both interfaces need not be configured twice. +func resolveShell[M adk.MessageType](cfg *BaseConfig[M]) adkfs.Shell { + if cfg.Shell != nil { + return cfg.Shell + } + if sh, ok := cfg.MemoryBackend.(adkfs.Shell); ok { + return sh + } + return nil +} + +func (m *middleware[M]) store() KVStore { + if m.cfg == nil { + return nil + } + return m.cfg.Store +} + +func (m *middleware[M]) lockTTL() time.Duration { + if m.cfg != nil && m.cfg.LockTTL > 0 { + return m.cfg.LockTTL + } + return defaultLockTTL +} + +func (m *middleware[M]) AfterAgent(ctx context.Context, _ *adk.TypedChatModelAgentState[M]) (context.Context, error) { + if m == nil || m.cfg == nil || m.sched == nil { + return ctx, nil + } + sessionID := m.sessionID + now := m.now() + touchTTL := m.sched.minInterval * 2 + if err := m.store().AddToSet(ctx, touchSetKey(m.resolvedMemoryDir), sessionID, now, touchTTL); err != nil { + m.onErr(ctx, OnErrorStageRecordTouch, err) + return ctx, nil + } + if err := m.maybeTrigger(ctx, sessionID, true); err != nil { + m.onErr(ctx, OnErrorStageRunDream, err) + } + return ctx, nil +} + +func (m *middleware[M]) maybeTrigger(ctx context.Context, currentSessionID string, excludeCurrent bool) error { + st, err := getScheduleState(ctx, m.store(), m.resolvedMemoryDir) + if err != nil { + return err + } + if st == nil { + st = &ScheduleState{} + } + now := m.now() + if st.NextCheckAt.After(now) { + return nil + } + since := st.LastConsolidatedAt + if !since.IsZero() && now.Sub(since) < m.sched.minInterval { + st.NextCheckAt = st.LastConsolidatedAt.Add(m.sched.minInterval) + return setScheduleState(ctx, m.store(), m.resolvedMemoryDir, st) + } + touchedSessions, err := m.store().ListSet(ctx, touchSetKey(m.resolvedMemoryDir), since) + if err != nil { + return err + } + // Filter in place: ListSet returns a freshly allocated slice the store no longer + // references, so reusing its backing array is safe. + filtered := touchedSessions[:0] + for _, sessionID := range touchedSessions { + if excludeCurrent && currentSessionID != "" && sessionID == currentSessionID { + continue + } + filtered = append(filtered, sessionID) + } + if len(filtered) < m.sched.minTouchedSession { + st.NextCheckAt = now.Add(m.sched.scanInterval) + return setScheduleState(ctx, m.store(), m.resolvedMemoryDir, st) + } + unlock, ok, err := m.store().AcquireLock(ctx, runLockKey(m.resolvedMemoryDir), m.lockTTL()) + if err != nil || !ok { + return err + } + + job := m.newJob(currentSessionID, filtered) + m.persistJob(ctx, job) + + // Detach from the request lifecycle (which ends when AfterAgent returns) while + // preserving context values such as loggers and trace spans. + runBaseCtx := withoutCancel(ctx) + runFn := func() { + defer func() { _ = unlock(runBaseCtx) }() + runErr := m.executeJob(runBaseCtx, job, currentSessionID, filtered) + if runErr != nil && job.Status != StatusCanceled { + st.ConsecutiveFailures++ + // After repeated failures, advance the window so the next run does not + // replay the same failing sessions forever. + if st.ConsecutiveFailures >= m.sched.maxConsecutiveFailures { + st.LastConsolidatedAt = m.now() + st.ConsecutiveFailures = 0 + } + st.NextCheckAt = m.now().Add(m.sched.scanInterval) + _ = setScheduleState(runBaseCtx, m.store(), m.resolvedMemoryDir, st) + return + } + if runErr != nil { // canceled: back off without counting as a failure + st.NextCheckAt = m.now().Add(m.sched.scanInterval) + _ = setScheduleState(runBaseCtx, m.store(), m.resolvedMemoryDir, st) + return + } + st.LastConsolidatedAt = m.now() + st.ConsecutiveFailures = 0 + st.NextCheckAt = st.LastConsolidatedAt.Add(m.sched.minInterval) + _ = setScheduleState(runBaseCtx, m.store(), m.resolvedMemoryDir, st) + // Touches consumed by this run are no longer needed; drop them so the set + // does not grow without bound. + _ = m.store().PruneSet(runBaseCtx, touchSetKey(m.resolvedMemoryDir), st.LastConsolidatedAt) + } + if m.sched.runInline { + runFn() + return nil + } + go runFn() + return nil +} + +// newJob builds a pending Job record for the given scope. +func (m *middleware[M]) newJob(sessionID string, sessionScope []string) *Job { + jobID := newJobID(m.resolvedMemoryDir + sessionID) + scope := sessionScope + if len(scope) == 0 && sessionID != "" { + scope = []string{sessionID} + } + return &Job{ + ID: jobID, + Status: StatusPending, + InputDir: m.resolvedMemoryDir, + OutputDir: m.resolvedOutputDir, + StagingDir: filepath.Join(m.stagingRoot, jobID), + SessionIDs: append([]string(nil), scope...), + CreatedAt: m.now(), + } +} + +func (m *middleware[M]) persistJob(ctx context.Context, job *Job) { + if m.store() == nil || job == nil { + return + } + if err := setJob(ctx, m.store(), job, jobTTL); err != nil { + m.onErr(ctx, OnErrorStagePersistJob, err) + } +} + +// executeJob runs one dream job through running -> completed/failed/canceled, +// persisting status transitions. It returns the consolidation error (nil on success). +func (m *middleware[M]) executeJob(ctx context.Context, job *Job, sessionID string, touchedSessions []string) error { + runCtx, cancel := context.WithCancel(ctx) + handle := &cancelHandle{cancel: cancel} + defer cancel() + registerCancel(job.ID, handle) + defer unregisterCancel(job.ID) + + job.Status = StatusRunning + m.persistJob(ctx, job) + + err := m.consolidate(runCtx, job, sessionID, touchedSessions) + + job.EndedAt = m.now() + switch { + case err != nil && m.wasCanceled(job.ID, runCtx, err): + job.Status = StatusCanceled + case err != nil: + job.Status = StatusFailed + job.ErrMsg = err.Error() + default: + job.Status = StatusCompleted + } + // Do not resurrect a job that was canceled out-of-band into a completed state. + if job.Status == StatusCompleted && m.store() != nil { + if latest, gerr := getJob(ctx, m.store(), job.ID); gerr == nil && latest != nil && latest.Status == StatusCanceled { + job.Status = StatusCanceled + } + } + m.persistJob(ctx, job) + return err +} + +func (m *middleware[M]) wasCanceled(jobID string, ctx context.Context, err error) bool { + if errors.Is(err, errDreamCanceled) { + return true + } + return errors.Is(jobCancelCause(jobID, ctx), errDreamCanceled) +} + +// consolidate seeds a staging working copy, runs the dream agent against it, then +// promotes the result to the output directory. The input directory is never modified. +func (m *middleware[M]) consolidate(ctx context.Context, job *Job, sessionID string, touchedSessions []string) error { + stagingFS, err := ainternal.NewFSBackend(m.cfg.MemoryBackend, ainternal.FSBackendConfig{ + BaseDir: job.StagingDir, + AllowLs: true, + NotFoundAsContent: true, + ErrorPrefix: "dream staging backend", + }) + if err != nil { + return err + } + + // Seed staging with a copy of the input directory so the model edits a working + // copy. Tolerate seed errors (e.g. an empty/new memory directory). + if seedErr := m.copyTree(ctx, m.inputFS, m.resolvedMemoryDir, stagingFS); seedErr != nil { + m.onErr(ctx, OnErrorStageSeedStaging, seedErr) + } + + fsHandler, err := fsmw.NewTyped[M](ctx, &fsmw.MiddlewareConfig{ + Backend: stagingFS, + GrepToolConfig: &fsmw.ToolConfig{Disable: true}, + }) + if err != nil { + return err + } + agent, err := m.newDreamAgent(ctx, fsHandler) + if err != nil { + return err + } + + prompt := buildConsolidationPrompt(job.StagingDir, touchedSessions, m.sessionSearchTool != nil) + searchSessionIDs := touchedSessions + if len(searchSessionIDs) == 0 && sessionID != "" { + searchSessionIDs = []string{sessionID} + } + runCtx := withDreamRunMeta(ctx, &dreamRunMeta{ + MemoryDirectory: m.resolvedMemoryDir, + SessionID: sessionID, + SearchSessionIDs: append([]string(nil), searchSessionIDs...), + }) + iter := agent.Run(runCtx, &adk.TypedAgentInput[M]{Messages: []M{makeUserMsg[M](prompt)}}) + if m.cfg.HandleIterator != nil { + teed, observedErr := m.teeIteratorErr(iter) + if err := m.cfg.HandleIterator(runCtx, teed); err != nil { + return err + } + if err := observedErr(); err != nil { + return err + } + } else if err := m.drainWithCancel(runCtx, job.ID, iter); err != nil { + return err + } + + return m.promote(ctx, stagingFS, job.StagingDir) +} + +// teeIteratorErr forwards every event from src to a fresh iterator handed to a custom +// HandleIterator, while recording the first event.Err seen on the stream. The +// returned func reports that error after the handler returns. This keeps failure +// detection correct regardless of whether the handler inspects event.Err itself. +func (m *middleware[M]) teeIteratorErr(src *adk.AsyncIterator[*adk.TypedAgentEvent[M]]) (*adk.AsyncIterator[*adk.TypedAgentEvent[M]], func() error) { + out, gen := adk.NewAsyncIteratorPair[*adk.TypedAgentEvent[M]]() + var mu sync.Mutex + var firstErr error + go func() { + defer gen.Close() + for { + ev, ok := src.Next() + if !ok { + return + } + if ev != nil && ev.Err != nil { + mu.Lock() + if firstErr == nil { + firstErr = ev.Err + } + mu.Unlock() + } + gen.Send(ev) + } + }() + return out, func() error { + mu.Lock() + defer mu.Unlock() + return firstErr + } +} + +// drainWithCancel drains the dream agent's event stream, aborting if the run is +// canceled (in-process via context, or cross-process via the job status in the store). +func (m *middleware[M]) drainWithCancel(ctx context.Context, jobID string, iter *adk.AsyncIterator[*adk.TypedAgentEvent[M]]) error { + i := 0 + for { + ev, ok := iter.Next() + if !ok { + return nil + } + if ev != nil && ev.Err != nil { + return ev.Err + } + if err := ctx.Err(); err != nil { + return jobCancelCause(jobID, ctx) + } + i++ + if m.store() != nil && i%cancelPollInterval == 0 { + if job, _ := getJob(ctx, m.store(), jobID); job != nil && job.Status == StatusCanceled { + return errDreamCanceled + } + } + } +} + +// promote copies the staged result to the output directory, holding the output +// directory lock so a concurrent dream (or, by convention, automemory extraction) +// does not interleave writes. Promotion is a non-atomic, best-effort copy. +func (m *middleware[M]) promote(ctx context.Context, stagingFS *ainternal.FSBackend, stagingDir string) error { + if m.store() != nil { + unlock, ok, err := m.store().AcquireLock(ctx, dirLockKey(m.resolvedOutputDir), m.lockTTL()) + if err != nil { + return err + } + if !ok { + return fmt.Errorf("dream: output directory %q is busy", m.resolvedOutputDir) + } + defer func() { _ = unlock(ctx) }() + } + if m.resolvedOutputDir == m.resolvedMemoryDir { + m.onErr(ctx, OnErrorStagePromote, errPromoteInPlaceBestEffort) + } + if err := m.copyTree(ctx, stagingFS, stagingDir, m.outputFS); err != nil { + return err + } + m.cleanupStaging(ctx, stagingDir) + return nil +} + +// cleanupStaging removes the staging directory using the resolved Shell. When no +// Shell is available the staging directory is left in place. +func (m *middleware[M]) cleanupStaging(ctx context.Context, stagingDir string) { + if m.shell == nil { + return + } + if _, err := m.shell.Execute(ctx, &adkfs.ExecuteRequest{ + Command: fmt.Sprintf("rm -rf %q", stagingDir), + }); err != nil { + m.onErr(ctx, OnErrorStageCleanup, err) + } +} + +// copyTree copies the memory files (markdown, per copyGlobPattern) from src +// (bounded at srcBase) to dst, preserving relative paths. Non-markdown files in the +// memory directory are left untouched. +func (m *middleware[M]) copyTree(ctx context.Context, src *ainternal.FSBackend, srcBase string, dst *ainternal.FSBackend) error { + files, err := src.GlobInfo(ctx, &adkfs.GlobInfoRequest{Path: srcBase, Pattern: copyGlobPattern}) + if err != nil { + return err + } + for _, fi := range files { + if fi.IsDir { + continue + } + rel, relErr := filepath.Rel(srcBase, fi.Path) + if relErr != nil { + rel = filepath.Base(fi.Path) + } + rel = filepath.ToSlash(rel) + fc, err := src.Read(ctx, &adkfs.ReadRequest{FilePath: rel}) + if err != nil { + return err + } + if fc == nil { + continue + } + if err := dst.Write(ctx, &adkfs.WriteRequest{FilePath: rel, Content: fc.Content}); err != nil { + return err + } + } + return nil +} + +func (m *middleware[M]) newDreamAgent(ctx context.Context, fsHandler adk.TypedChatModelAgentMiddleware[M]) (*adk.TypedChatModelAgent[M], error) { + tools := make([]tool.BaseTool, 0, 1) + if m.sessionSearchTool != nil { + tools = append(tools, m.sessionSearchTool) + } + maxIterations := m.cfg.MaxIterations + if maxIterations <= 0 { + maxIterations = defaultMaxIterations + } + agent, err := adk.NewTypedChatModelAgent(ctx, &adk.TypedChatModelAgentConfig[M]{ + Name: "automemory_dream", + Description: "Internal auto dream consolidation agent", + Model: m.cfg.Model, + Handlers: []adk.TypedChatModelAgentMiddleware[M]{fsHandler}, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: tools, + ToolCallMiddlewares: []compose.ToolMiddleware{m.toolErrorRecoveryMiddleware(ctx)}, + }}, + MaxIterations: maxIterations, + }) + if err != nil { + return nil, fmt.Errorf("auto dream create agent: %w", err) + } + return agent, nil +} + +// toolErrorRecoveryMiddleware turns a failed tool call into a message fed back to the +// model, so a recoverable failure (e.g. an edit_file whose old_string no longer +// matches) lets the agent retry or switch tools rather than aborting the whole run. +// The error is still reported through OnError for observability. Runtime errors that +// the model cannot act on (model-call failures, context cancellation) are not tool +// errors and are unaffected — they continue to fail the run. +func (m *middleware[M]) toolErrorRecoveryMiddleware(ctx context.Context) compose.ToolMiddleware { + recover := func(name string, err error) string { + m.onErr(ctx, OnErrorStageToolCall, fmt.Errorf("tool %q failed: %w", name, err)) + return fmt.Sprintf("Tool %q failed: %v\nThis is not fatal. Re-check your arguments "+ + "(for edits, read_file the target again and correct old_string), then retry or use a "+ + "different tool. Do not repeat the same failing call.", name, err) + } + return compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, in *compose.ToolInput) (*compose.ToolOutput, error) { + out, err := next(ctx, in) + if err != nil { + return &compose.ToolOutput{Result: recover(in.Name, err)}, nil + } + return out, nil + } + }, + EnhancedInvokable: func(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint { + return func(ctx context.Context, in *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + out, err := next(ctx, in) + if err != nil { + return &compose.EnhancedInvokableToolOutput{ + Result: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: recover(in.Name, err)}}, + }, + }, nil + } + return out, nil + } + }, + } +} + +func (m *middleware[M]) onErr(ctx context.Context, stage ErrorStage, err error) { + if err == nil || m == nil || m.cfg == nil || m.cfg.OnError == nil { + return + } + m.cfg.OnError(ctx, stage, err) +} + +func makeUserMsg[M adk.MessageType](text string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(text)).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(text)).(M) + default: + panic("unreachable") + } +} diff --git a/adk/middlewares/automemory/dream/prompt.go b/adk/middlewares/automemory/dream/prompt.go new file mode 100644 index 000000000..075a80134 --- /dev/null +++ b/adk/middlewares/automemory/dream/prompt.go @@ -0,0 +1,141 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "fmt" + "strings" + + "github.com/cloudwego/eino/adk/internal" +) + +func buildConsolidationPrompt(memoryRoot string, touchedSessions []string, includeSessionSearch bool) string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: buildConsolidationPromptEnglish(memoryRoot, touchedSessions, includeSessionSearch), + Chinese: buildConsolidationPromptChinese(memoryRoot, touchedSessions, includeSessionSearch), + }) +} + +func buildConsolidationPromptEnglish(memoryRoot string, touchedSessions []string, includeSessionSearch bool) string { + extra := "" + if len(touchedSessions) > 0 { + extra = fmt.Sprintf("\n\nSessions since last consolidation (%d):\n%s", len(touchedSessions), bulletList(touchedSessions)) + } + sessionSearchSection := "" + if includeSessionSearch { + sessionSearchSection = ` + +## Optional session search + +- Use grep_session_history with narrow terms when you already suspect something matters +- It searches only the session histories included in this dream run +- Do not exhaustively scan session history; use it only to confirm details` + } + return fmt.Sprintf(`# Dream: Memory Consolidation + +You are performing a dream: a reflective pass over persistent memory files. Synthesize what was learned recently into durable, well-organized memory so future sessions can orient quickly. + +Memory directory: %s + +This directory is a working copy of the memory. Edit it freely: your result is promoted to the live memory location only after this run completes successfully. Operate only within this directory. + +## Phase 1 - Orient +- Use ls/glob to inspect the memory directory +- Read MEMORY.md first to understand the current index +- Skim existing topic files before creating new ones so you improve or merge instead of duplicating%s + +## Phase 2 - Gather signal +- Focus on durable information that has emerged across recent sessions +- Prefer updating an existing topic file over creating a near-duplicate +- Convert relative time references into absolute dates when they matter +- Remove or correct stale facts at the source + +## Phase 3 - Consolidate +- Keep each memory file focused on one topic +- Use read_file before write_file/edit_file for every file you plan to touch +- Write only inside the memory directory +- Do not investigate the codebase outside memory files and the current session history during this run + +## Phase 4 - Prune and index +- Keep MEMORY.md concise; it is an index, not the full memory body +- Ensure new or updated topic files are reflected in MEMORY.md +- Remove stale or superseded pointers from MEMORY.md + +Return a brief summary of what you consolidated, updated, or pruned. If nothing changed, say so.%s`, memoryRoot, sessionSearchSection, extra) +} + +func buildConsolidationPromptChinese(memoryRoot string, touchedSessions []string, includeSessionSearch bool) string { + extra := "" + if len(touchedSessions) > 0 { + extra = fmt.Sprintf("\n\n自上次 consolidation 以来触达过的 sessions（%d）：\n%s", len(touchedSessions), bulletList(touchedSessions)) + } + sessionSearchSection := "" + if includeSessionSearch { + sessionSearchSection = ` + +## 可选的 session 搜索 + +- 当你已经怀疑某条信息重要时，再用 grep_session_history 做精确搜索 +- 它只会搜索本次 dream 运行范围内包含的 session 历史 +- 不要穷举扫描 session 历史，只在需要核实细节时使用` + } + return fmt.Sprintf(`# Dream：记忆整理 + +你正在执行一次 dream：对持久化记忆文件做反思式整理。请把最近学到的内容沉淀成稳定、清晰且结构化的长期记忆，帮助未来会话快速建立上下文。 + +记忆目录：%s + +该目录是记忆的工作副本，可放心修改：只有在本次运行成功完成后，整理结果才会被提升（promote）到真正的记忆位置。请只在该目录内操作。 + +## 阶段 1 - 建立整体认识 +- 使用 ls/glob 查看记忆目录 +- 先阅读 MEMORY.md，理解当前索引结构 +- 在创建新主题文件前，先浏览现有主题文件，优先改进或合并，而不是重复创建%s + +## 阶段 2 - 收集有效信号 +- 关注在最近多个 session 中沉淀下来的长期有效信息 +- 优先更新已有主题文件，而不是创建内容接近的重复文件 +- 当相对时间表述会影响理解时，将其转换为绝对日期 +- 在源头处删除或修正过时事实 + +## 阶段 3 - 整理与归并 +- 让每个记忆文件只聚焦一个主题 +- 对每个计划修改的文件，都先 read_file，再 write_file/edit_file +- 只在记忆目录内写入 +- 本次运行中，不要调查记忆文件与当前 session 历史之外的代码库内容 + +## 阶段 4 - 修剪与更新索引 +- 保持 MEMORY.md 简洁；它是索引，不是完整记忆正文 +- 确保新增或更新过的主题文件都同步反映到 MEMORY.md +- 从 MEMORY.md 中移除陈旧或已被替代的索引项 + +请简要总结你本次 consolidation、更新或修剪了什么；如果没有任何变更，也请明确说明。%s`, memoryRoot, sessionSearchSection, extra) +} + +func bulletList(items []string) string { + if len(items) == 0 { + return "" + } + var b strings.Builder + b.WriteString("- ") + b.WriteString(items[0]) + for i := 1; i < len(items); i++ { + b.WriteString("\n- ") + b.WriteString(items[i]) + } + return b.String() +} diff --git a/adk/middlewares/automemory/dream/session.go b/adk/middlewares/automemory/dream/session.go new file mode 100644 index 000000000..32c5f952a --- /dev/null +++ b/adk/middlewares/automemory/dream/session.go @@ -0,0 +1,177 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "fmt" + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + toolutils "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/schema" +) + +type grepSessionHistoryInput struct { + Query string `json:"query" jsonschema:"required,description=the narrow term to search in current session history"` + Limit int `json:"limit,omitempty" jsonschema:"description=maximum number of matching lines to return"` +} + +type dreamRunMeta struct { + MemoryDirectory string + SessionID string + SearchSessionIDs []string +} + +type dreamRunMetaKey struct{} + +func withDreamRunMeta(ctx context.Context, meta *dreamRunMeta) context.Context { + return context.WithValue(ctx, dreamRunMetaKey{}, meta) +} + +func getDreamRunMeta(ctx context.Context) *dreamRunMeta { + if v := ctx.Value(dreamRunMetaKey{}); v != nil { + if meta, ok := v.(*dreamRunMeta); ok { + return meta + } + } + return nil +} + +func newSessionHistoryGrepTool[M adk.MessageType](store adk.SessionEventStore[M]) (tool.BaseTool, error) { + if store == nil { + return nil, nil + } + + t, err := toolutils.InferTool("grep_session_history", internal.SelectPrompt(internal.I18nPrompts{ + English: "Search the session histories included in the current dream run with a narrow query and return matching lines.", + Chinese: "在当前 dream 运行范围内的会话历史中按精确关键词搜索，并返回匹配行。", + }), func(ctx context.Context, input grepSessionHistoryInput) (string, error) { + meta := getDreamRunMeta(ctx) + if meta == nil { + return "", fmt.Errorf("grep_session_history: missing dream run metadata") + } + sessionIDs := resolveSearchSessionIDs(meta) + if len(sessionIDs) == 0 { + return "", fmt.Errorf("grep_session_history: no searchable sessions in current dream run") + } + query := strings.TrimSpace(input.Query) + if query == "" { + return "", fmt.Errorf("grep_session_history: empty query") + } + limit := input.Limit + if limit <= 0 { + limit = 50 + } + + pageSize := limit + if pageSize < 100 { + pageSize = 100 + } + + var ( + after string + found []string + ) + includeSessionPrefix := len(sessionIDs) > 1 + for _, sessionID := range sessionIDs { + after = "" + for len(found) < limit { + result, err := store.LoadEvents(ctx, sessionID, &adk.LoadSessionEventsRequest{ + After: after, + Limit: pageSize, + Reverse: true, + Kinds: []adk.SessionEventKind{adk.SessionEventMessage}, + }) + if err != nil { + return "", err + } + if result == nil || len(result.Events) == 0 { + break + } + for _, ev := range result.Events { + found = appendMatchingSessionHistoryLines(found, sessionID, sessionEventMessageString(ev.Message), query, limit, includeSessionPrefix) + if len(found) >= limit { + break + } + } + if result.Next == "" { + break + } + after = result.Next + } + if len(found) >= limit { + break + } + } + + return strings.Join(found, "\n"), nil + }) + if err != nil { + return nil, err + } + + return t, nil +} + +func resolveSearchSessionIDs(meta *dreamRunMeta) []string { + if meta == nil { + return nil + } + if len(meta.SearchSessionIDs) > 0 { + return meta.SearchSessionIDs + } + if meta.SessionID != "" { + return []string{meta.SessionID} + } + return nil +} + +func appendMatchingSessionHistoryLines(dst []string, sessionID, message, query string, limit int, includeSessionPrefix bool) []string { + needle := strings.ToLower(query) + for _, line := range strings.Split(message, "\n") { + if strings.Contains(strings.ToLower(line), needle) { + if includeSessionPrefix { + line = fmt.Sprintf("[%s] %s", sessionID, line) + } + dst = append(dst, line) + if len(dst) >= limit { + return dst + } + } + } + return dst +} + +func sessionEventMessageString[M adk.MessageType](msg M) string { + switch m := any(msg).(type) { + case *schema.Message: + if m == nil { + return "" + } + return m.String() + case *schema.AgenticMessage: + if m == nil { + return "" + } + return m.String() + default: + return "" + } +} diff --git a/adk/middlewares/automemory/dream/session_test.go b/adk/middlewares/automemory/dream/session_test.go new file mode 100644 index 000000000..bb8f66818 --- /dev/null +++ b/adk/middlewares/automemory/dream/session_test.go @@ -0,0 +1,111 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + adksession "github.com/cloudwego/eino/adk/session" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func TestNewSessionHistoryGrepTool(t *testing.T) { + ctx := context.Background() + store := adksession.NewInMemoryStore[*schema.Message](nil) + sessionID := "session-1" + + appendEvent := func(eventID string, msg *schema.Message) { + err := store.AppendEvents(ctx, sessionID, []*adk.SessionEvent[*schema.Message]{{ + EventID: eventID, + Kind: adk.SessionEventMessage, + Message: msg, + }}) + require.NoError(t, err) + } + + appendEvent("e1", schema.UserMessage("hello there")) + appendEvent("e2", schema.AssistantMessage("build failure: missing dependency", nil)) + appendEvent("e3", schema.ToolMessage("Build Failure: retry later", "call-1")) + + bt, err := newSessionHistoryGrepTool[*schema.Message](store) + require.NoError(t, err) + + result, err := bt.(tool.InvokableTool).InvokableRun( + withDreamRunMeta(ctx, &dreamRunMeta{SessionID: sessionID, SearchSessionIDs: []string{sessionID}}), + `{"query":"build failure","limit":2}`, + ) + require.NoError(t, err) + require.Equal(t, "tool: Build Failure: retry later\nassistant: build failure: missing dependency", result) +} + +func TestNewSessionHistoryGrepTool_SearchesRunScopedSessions(t *testing.T) { + ctx := context.Background() + store := adksession.NewInMemoryStore[*schema.Message](nil) + + appendEvent := func(sessionID, eventID string, msg *schema.Message) { + err := store.AppendEvents(ctx, sessionID, []*adk.SessionEvent[*schema.Message]{{ + EventID: eventID, + Kind: adk.SessionEventMessage, + Message: msg, + }}) + require.NoError(t, err) + } + + appendEvent("session-a", "a1", schema.AssistantMessage("build failure: missing dependency", nil)) + appendEvent("session-b", "b1", schema.ToolMessage("build failure: retry later", "call-1")) + appendEvent("session-c", "c1", schema.AssistantMessage("build failure: should not be searched", nil)) + + bt, err := newSessionHistoryGrepTool[*schema.Message](store) + require.NoError(t, err) + + result, err := bt.(tool.InvokableTool).InvokableRun( + withDreamRunMeta(ctx, &dreamRunMeta{ + SessionID: "session-c", + SearchSessionIDs: []string{"session-a", "session-b"}, + }), + `{"query":"build failure","limit":5}`, + ) + require.NoError(t, err) + require.Contains(t, result, "[session-a] assistant: build failure: missing dependency") + require.Contains(t, result, "[session-b] tool: build failure: retry later") + require.NotContains(t, result, "should not be searched") +} + +func TestNewSessionHistoryGrepTool_InfoUsesChineseDescription(t *testing.T) { + require.NoError(t, adk.SetLanguage(adk.LanguageChinese)) + defer func() { + require.NoError(t, adk.SetLanguage(adk.LanguageEnglish)) + }() + + bt, err := newSessionHistoryGrepTool[*schema.Message](adksession.NewInMemoryStore[*schema.Message](nil)) + require.NoError(t, err) + + info, err := bt.Info(context.Background()) + require.NoError(t, err) + require.Contains(t, info.Desc, "在当前 dream 运行范围内的会话历史中按精确关键词搜索") +} + +func TestNewSessionHistoryGrepTool_AllowsNilStore(t *testing.T) { + bt, err := newSessionHistoryGrepTool[*schema.Message](nil) + require.NoError(t, err) + require.Nil(t, bt) +} diff --git a/adk/middlewares/automemory/dream/store.go b/adk/middlewares/automemory/dream/store.go new file mode 100644 index 000000000..28be51635 --- /dev/null +++ b/adk/middlewares/automemory/dream/store.go @@ -0,0 +1,266 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package dream + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "sync" + "time" +) + +// KVStore is the Redis-shaped coordination backend dream relies on. It persists +// scheduling state, the touched-session set, run locks, and dream job records. +// +// The shape deliberately mirrors a key-value store with locks so it maps cleanly +// onto Redis (AcquireLock -> SET NX PX, Get -> GET, Set -> SET PX, Del -> DEL) and +// onto a sorted set for the touch operations (AddToSet -> ZADD, ListSet -> ZRANGEBYSCORE, +// PruneSet -> ZREMRANGEBYSCORE). +// +// The in-process NewLocalKVStore is the default, but it is single-process only: +// when the dream middleware is constructed per session (the common server pattern), +// each instance gets its own store, touched-session counts never reach the trigger +// threshold, and dreams never fire. Production deployments MUST inject a shared, +// durable KVStore whose AcquireLock is atomic across processes. +type KVStore interface { + // AcquireLock tries to acquire a lock for key. When ok==true it returns an + // unlock function that must be called exactly once. ok==false means another + // holder owns the lock. + AcquireLock(ctx context.Context, key string, ttl time.Duration) (unlock func(context.Context) error, ok bool, err error) + + // Get returns the value for key. When the key does not exist, ok is false. + Get(ctx context.Context, key string) (value []byte, ok bool, err error) + + // Set stores value for key. ttl<=0 means no expiration. + Set(ctx context.Context, key string, value []byte, ttl time.Duration) error + + // Del removes key. Removing a missing key is a no-op. + Del(ctx context.Context, key string) error + + // AddToSet adds member to the set at key, scored by at. ttl<=0 means no + // expiration on the set. + AddToSet(ctx context.Context, key, member string, at time.Time, ttl time.Duration) error + + // ListSet returns distinct members of the set at key scored strictly after since. + ListSet(ctx context.Context, key string, since time.Time) ([]string, error) + + // PruneSet removes members of the set at key scored at or before before. + PruneSet(ctx context.Context, key string, before time.Time) error +} + +// ScheduleState stores per-memory-directory scheduling state. +type ScheduleState struct { + // LastConsolidatedAt is the completion time of the last successful run. + LastConsolidatedAt time.Time `json:"last_consolidated_at"` + + // NextCheckAt is the next time the middleware should re-check this directory. + NextCheckAt time.Time `json:"next_check_at"` + + // ConsecutiveFailures counts dream runs that have failed since the last + // success. It is reset to zero on a successful run. + ConsecutiveFailures int `json:"consecutive_failures"` +} + +// Key derivation: all keys are scoped to the resolved memory directory so multiple +// memory directories sharing one KVStore do not collide. + +func dirHash(memoryDir string) string { + sum := sha256.Sum256([]byte(memoryDir)) + return hex.EncodeToString(sum[:8]) +} + +func scheduleKey(memoryDir string) string { return "dream::" + dirHash(memoryDir) + "::schedule" } +func touchSetKey(memoryDir string) string { return "dream::" + dirHash(memoryDir) + "::touch" } +func runLockKey(memoryDir string) string { return "dream::" + dirHash(memoryDir) + "::lock" } +func dirLockKey(dir string) string { return "dream::" + dirHash(dir) + "::dirlock" } +func jobKey(jobID string) string { return "dream::job::" + jobID } + +func getScheduleState(ctx context.Context, store KVStore, memoryDir string) (*ScheduleState, error) { + raw, ok, err := store.Get(ctx, scheduleKey(memoryDir)) + if err != nil || !ok { + return nil, err + } + var st ScheduleState + if err := json.Unmarshal(raw, &st); err != nil { + return nil, err + } + return &st, nil +} + +func setScheduleState(ctx context.Context, store KVStore, memoryDir string, st *ScheduleState) error { + if st == nil { + return store.Del(ctx, scheduleKey(memoryDir)) + } + raw, err := json.Marshal(st) + if err != nil { + return err + } + return store.Set(ctx, scheduleKey(memoryDir), raw, 0) +} + +func getJob(ctx context.Context, store KVStore, jobID string) (*Job, error) { + raw, ok, err := store.Get(ctx, jobKey(jobID)) + if err != nil || !ok { + return nil, err + } + var job Job + if err := json.Unmarshal(raw, &job); err != nil { + return nil, err + } + return &job, nil +} + +func setJob(ctx context.Context, store KVStore, job *Job, ttl time.Duration) error { + raw, err := json.Marshal(job) + if err != nil { + return err + } + return store.Set(ctx, jobKey(job.ID), raw, ttl) +} + +// localKVStore is the default in-process KVStore. It is single-process only and is +// suitable for tests and single-instance deployments. +type localKVStore struct { + mu sync.Mutex + kv map[string]localKVValue + locks map[string]localKVLock + sets map[string]map[string]time.Time +} + +type localKVValue struct { + value []byte + expiry time.Time +} + +type localKVLock struct { + token string + expiry time.Time +} + +// NewLocalKVStore returns an in-process KVStore. +// It is suitable for tests and single-process use only. +func NewLocalKVStore() KVStore { + return &localKVStore{ + kv: make(map[string]localKVValue), + locks: make(map[string]localKVLock), + sets: make(map[string]map[string]time.Time), + } +} + +func (s *localKVStore) AcquireLock(_ context.Context, key string, ttl time.Duration) (func(context.Context) error, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + now := time.Now() + if l, ok := s.locks[key]; ok && now.Before(l.expiry) { + return nil, false, nil + } + token := randToken() + s.locks[key] = localKVLock{token: token, expiry: now.Add(ttl)} + return func(context.Context) error { + s.mu.Lock() + defer s.mu.Unlock() + l, ok := s.locks[key] + if !ok { + return nil + } + if l.token != token { + return fmt.Errorf("lock token mismatch") + } + delete(s.locks, key) + return nil + }, true, nil +} + +func (s *localKVStore) Get(_ context.Context, key string) ([]byte, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.kv[key] + if !ok { + return nil, false, nil + } + if !v.expiry.IsZero() && time.Now().After(v.expiry) { + delete(s.kv, key) + return nil, false, nil + } + return append([]byte(nil), v.value...), true, nil +} + +func (s *localKVStore) Set(_ context.Context, key string, value []byte, ttl time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + var expiry time.Time + if ttl > 0 { + expiry = time.Now().Add(ttl) + } + s.kv[key] = localKVValue{value: append([]byte(nil), value...), expiry: expiry} + return nil +} + +func (s *localKVStore) Del(_ context.Context, key string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.kv, key) + return nil +} + +func (s *localKVStore) AddToSet(_ context.Context, key, member string, at time.Time, _ time.Duration) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.sets[key] == nil { + s.sets[key] = make(map[string]time.Time) + } + s.sets[key][member] = at + return nil +} + +func (s *localKVStore) ListSet(_ context.Context, key string, since time.Time) ([]string, error) { + s.mu.Lock() + defer s.mu.Unlock() + items := s.sets[key] + if len(items) == 0 { + return nil, nil + } + out := make([]string, 0, len(items)) + for member, at := range items { + if at.After(since) { + out = append(out, member) + } + } + return out, nil +} + +func (s *localKVStore) PruneSet(_ context.Context, key string, before time.Time) error { + s.mu.Lock() + defer s.mu.Unlock() + items := s.sets[key] + for member, at := range items { + if !at.After(before) { + delete(items, member) + } + } + return nil +} + +func randToken() string { + var b [8]byte + _, _ = rand.Read(b[:]) + return hex.EncodeToString(b[:]) +} diff --git a/adk/middlewares/automemory/inmemory_backend.go b/adk/middlewares/automemory/inmemory_backend.go new file mode 100644 index 000000000..9b775b139 --- /dev/null +++ b/adk/middlewares/automemory/inmemory_backend.go @@ -0,0 +1,210 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +import ( + "context" + "fmt" + "path/filepath" + "sort" + "strings" + "sync" + "time" + + "github.com/bmatcuk/doublestar/v4" +) + +type memFile struct { + content string + modifiedAt time.Time +} + +// InMemoryBackend is a simple in-memory Backend implementation intended for tests +// and demos. Paths are treated as filesystem-like, and should be absolute. +type InMemoryBackend struct { + mu sync.RWMutex + files map[string]*memFile +} + +// NewInMemoryBackend returns an empty in-memory Backend implementation. +func NewInMemoryBackend() *InMemoryBackend { + return &InMemoryBackend{ + files: make(map[string]*memFile), + } +} + +func (b *InMemoryBackend) put(path string, content string, modifiedAt time.Time) { + b.mu.Lock() + defer b.mu.Unlock() + b.files[filepath.Clean(path)] = &memFile{content: content, modifiedAt: modifiedAt} +} + +func (b *InMemoryBackend) Write(_ context.Context, req *WriteRequest) error { + if req == nil || req.FilePath == "" { + return fmt.Errorf("write: invalid request") + } + // Default to full replace. + b.put(req.FilePath, req.Content, time.Now()) + return nil +} + +func (b *InMemoryBackend) Edit(_ context.Context, req *EditRequest) error { + if req == nil || req.FilePath == "" { + return fmt.Errorf("edit: invalid request") + } + b.mu.Lock() + defer b.mu.Unlock() + + path := filepath.Clean(req.FilePath) + f, ok := b.files[path] + if !ok { + return fmt.Errorf("file not found: %s", path) + } + + if req.OldString == "" { + return fmt.Errorf("edit: old string must be non-empty") + } + if req.OldString == req.NewString { + return fmt.Errorf("edit: new string must differ from old string") + } + + out := f.content + if req.ReplaceAll { + out = strings.ReplaceAll(out, req.OldString, req.NewString) + } else { + if strings.Count(out, req.OldString) != 1 { + return fmt.Errorf("edit: old string must appear exactly once when ReplaceAll is false") + } + out = strings.Replace(out, req.OldString, req.NewString, 1) + } + f.content = out + f.modifiedAt = time.Now() + return nil +} + +func (b *InMemoryBackend) Read(_ context.Context, req *ReadRequest) (*FileContent, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + if req == nil || req.FilePath == "" { + return nil, fmt.Errorf("read: invalid request") + } + path := filepath.Clean(req.FilePath) + f, ok := b.files[path] + if !ok { + return nil, fmt.Errorf("file not found: %s", path) + } + + offset := req.Offset - 1 + if offset < 0 { + offset = 0 + } + limit := req.Limit + + content := f.content + if offset == 0 && limit <= 0 { + return &FileContent{Content: content}, nil + } + + start := 0 + for i := 0; i < offset; i++ { + idx := strings.IndexByte(content[start:], '\n') + if idx == -1 { + return &FileContent{Content: ""}, nil + } + start += idx + 1 + } + + if limit <= 0 { + return &FileContent{Content: content[start:]}, nil + } + + end := start + for i := 0; i < limit; i++ { + idx := strings.IndexByte(content[end:], '\n') + if idx == -1 { + return &FileContent{Content: content[start:]}, nil + } + end += idx + 1 + } + + // Trim trailing newline. + return &FileContent{Content: strings.TrimSuffix(content[start:end], "\n")}, nil +} + +func (b *InMemoryBackend) GlobInfo(_ context.Context, req *GlobInfoRequest) ([]FileInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + if req == nil || req.Pattern == "" { + return nil, fmt.Errorf("glob: invalid request") + } + base := filepath.Clean(req.Path) + if base == "." { + base = "" + } + + type item struct { + fi FileInfo + t time.Time + } + var out []item + + for p, f := range b.files { + if base != "" { + // Require p under base. + if p != base && !strings.HasPrefix(p, base+string(filepath.Separator)) { + continue + } + } + + rel := p + if base != "" { + rel = strings.TrimPrefix(p, base+string(filepath.Separator)) + if rel == p { + rel = strings.TrimPrefix(p, base) + rel = strings.TrimPrefix(rel, string(filepath.Separator)) + } + } + rel = filepath.ToSlash(rel) + + ok, err := doublestar.Match(req.Pattern, rel) + if err != nil { + return nil, err + } + if !ok { + continue + } + + out = append(out, item{ + fi: FileInfo{ + Path: p, + IsDir: false, + Size: int64(len(f.content)), + ModifiedAt: f.modifiedAt.Format(time.RFC3339Nano), + }, + t: f.modifiedAt, + }) + } + + sort.Slice(out, func(i, j int) bool { return out[i].t.After(out[j].t) }) + ret := make([]FileInfo, 0, len(out)) + for _, it := range out { + ret = append(ret, it.fi) + } + return ret, nil +} diff --git a/adk/middlewares/automemory/internal/backend.go b/adk/middlewares/automemory/internal/backend.go new file mode 100644 index 000000000..6f866764e --- /dev/null +++ b/adk/middlewares/automemory/internal/backend.go @@ -0,0 +1,235 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package internal contains bounded filesystem adapters used by automemory +// middleware implementations. +package internal + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + adkfs "github.com/cloudwego/eino/adk/filesystem" +) + +type Backend interface { + Read(ctx context.Context, req *adkfs.ReadRequest) (*adkfs.FileContent, error) + GlobInfo(ctx context.Context, req *adkfs.GlobInfoRequest) ([]adkfs.FileInfo, error) + Write(ctx context.Context, req *adkfs.WriteRequest) error + Edit(ctx context.Context, req *adkfs.EditRequest) error +} + +type FSBackendConfig struct { + BaseDir string + AllowLs bool + AllowGrep bool + NotFoundAsContent bool + ErrorPrefix string +} + +type FSBackend struct { + backend Backend + baseClean string + allowLs bool + allowGrep bool + notFoundAsContent bool + errorPrefix string +} + +// ResolveMemoryDir returns the cleaned absolute path for a memory directory. +func ResolveMemoryDir(dir string) (string, error) { + abs, err := filepath.Abs(dir) + if err != nil { + return "", err + } + return filepath.Clean(abs), nil +} + +// NewFSBackend wraps a Backend with path-bounding and optional tool behaviors. +func NewFSBackend(backend Backend, cfg FSBackendConfig) (*FSBackend, error) { + if backend == nil { + return nil, fmt.Errorf("%s: nil backend", prefixOrDefault(cfg.ErrorPrefix)) + } + if cfg.BaseDir == "" { + return nil, fmt.Errorf("%s: empty base dir", prefixOrDefault(cfg.ErrorPrefix)) + } + baseClean, err := ResolveMemoryDir(cfg.BaseDir) + if err != nil { + return nil, fmt.Errorf("%s: resolve base dir: %w", prefixOrDefault(cfg.ErrorPrefix), err) + } + return &FSBackend{ + backend: backend, + baseClean: baseClean, + allowLs: cfg.AllowLs, + allowGrep: cfg.AllowGrep, + notFoundAsContent: cfg.NotFoundAsContent, + errorPrefix: prefixOrDefault(cfg.ErrorPrefix), + }, nil +} + +func prefixOrDefault(prefix string) string { + if prefix == "" { + return "fs backend" + } + return prefix +} + +func isFileNotFoundErr(err error) bool { + if err == nil { + return false + } + if os.IsNotExist(err) { + return true + } + msg := strings.ToLower(err.Error()) + return strings.Contains(msg, "file not found") || strings.Contains(msg, "no such file or directory") +} + +func (f *FSBackend) resolveFilePath(p string) (string, error) { + if p == "" { + return "", fmt.Errorf("%s: empty path", f.errorPrefix) + } + if !filepath.IsAbs(p) { + p = filepath.Join(f.baseClean, p) + } + p = filepath.Clean(p) + if p != f.baseClean && !strings.HasPrefix(p, f.baseClean+string(filepath.Separator)) { + return "", fmt.Errorf("%s: path out of bounds: %s", f.errorPrefix, p) + } + return p, nil +} + +func (f *FSBackend) resolveDirPath(p string) (string, error) { + if p == "" { + return f.baseClean, nil + } + if !filepath.IsAbs(p) { + p = filepath.Join(f.baseClean, p) + } + p = filepath.Clean(p) + if p != f.baseClean && !strings.HasPrefix(p, f.baseClean+string(filepath.Separator)) { + return "", fmt.Errorf("%s: dir out of bounds: %s", f.errorPrefix, p) + } + return p, nil +} + +func (f *FSBackend) Read(ctx context.Context, req *adkfs.ReadRequest) (*adkfs.FileContent, error) { + if req == nil { + return nil, fmt.Errorf("read: invalid request") + } + fp, err := f.resolveFilePath(req.FilePath) + if err != nil { + return nil, err + } + n := *req + n.FilePath = fp + content, err := f.backend.Read(ctx, &n) + if err != nil { + if f.notFoundAsContent && isFileNotFoundErr(err) { + return &adkfs.FileContent{Content: fmt.Sprintf("File not found: %s", fp)}, nil + } + return nil, err + } + return content, nil +} + +func (f *FSBackend) Write(ctx context.Context, req *adkfs.WriteRequest) error { + if req == nil { + return fmt.Errorf("write: invalid request") + } + fp, err := f.resolveFilePath(req.FilePath) + if err != nil { + return err + } + n := *req + n.FilePath = fp + return f.backend.Write(ctx, &n) +} + +func (f *FSBackend) Edit(ctx context.Context, req *adkfs.EditRequest) error { + if req == nil { + return fmt.Errorf("edit: invalid request") + } + fp, err := f.resolveFilePath(req.FilePath) + if err != nil { + return err + } + n := *req + n.FilePath = fp + return f.backend.Edit(ctx, &n) +} + +func (f *FSBackend) GlobInfo(ctx context.Context, req *adkfs.GlobInfoRequest) ([]adkfs.FileInfo, error) { + if req == nil || req.Pattern == "" { + return nil, fmt.Errorf("glob: invalid request") + } + pathAbs, err := f.resolveDirPath(req.Path) + if err != nil { + return nil, err + } + pattern := req.Pattern + if filepath.IsAbs(pattern) { + cp := filepath.Clean(pattern) + if cp == pathAbs { + pattern = "." + } else if strings.HasPrefix(cp, pathAbs+string(filepath.Separator)) { + rel, rerr := filepath.Rel(pathAbs, cp) + if rerr != nil { + return nil, rerr + } + pattern = filepath.ToSlash(rel) + } else if strings.HasPrefix(cp, f.baseClean+string(filepath.Separator)) { + rel, rerr := filepath.Rel(f.baseClean, cp) + if rerr != nil { + return nil, rerr + } + pattern = filepath.ToSlash(rel) + pathAbs = f.baseClean + } else { + return nil, fmt.Errorf("%s: glob pattern out of bounds: %s", f.errorPrefix, cp) + } + } else { + pattern = filepath.ToSlash(pattern) + } + n := *req + n.Path = pathAbs + n.Pattern = pattern + return f.backend.GlobInfo(ctx, &n) +} + +func (f *FSBackend) LsInfo(ctx context.Context, req *adkfs.LsInfoRequest) ([]adkfs.FileInfo, error) { + if !f.allowLs { + return nil, fmt.Errorf("ls: disabled") + } + if req == nil { + return nil, fmt.Errorf("ls: invalid request") + } + base, err := f.resolveDirPath(req.Path) + if err != nil { + return nil, err + } + return f.GlobInfo(ctx, &adkfs.GlobInfoRequest{Path: base, Pattern: "*"}) +} + +func (f *FSBackend) GrepRaw(context.Context, *adkfs.GrepRequest) ([]adkfs.GrepMatch, error) { + if !f.allowGrep { + return nil, fmt.Errorf("grep: disabled") + } + return nil, fmt.Errorf("grep: not implemented") +} diff --git a/adk/middlewares/automemory/local_backend.go b/adk/middlewares/automemory/local_backend.go new file mode 100644 index 000000000..d437caaef --- /dev/null +++ b/adk/middlewares/automemory/local_backend.go @@ -0,0 +1,192 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +import ( + "context" + "fmt" + "io/fs" + "os" + "path/filepath" + "sort" + "strings" + "time" + + "github.com/bmatcuk/doublestar/v4" +) + +// LocalBackend implements Backend on the local OS filesystem. +// It is intentionally minimal (Read + GlobInfo) to match the "方案一" storage abstraction. +type LocalBackend struct{} + +// NewLocalBackend returns a filesystem-backed Backend implementation. +func NewLocalBackend() *LocalBackend { + return &LocalBackend{} +} + +func (b *LocalBackend) Read(_ context.Context, req *ReadRequest) (*FileContent, error) { + if req == nil || req.FilePath == "" { + return nil, fmt.Errorf("read: invalid request") + } + + raw, err := os.ReadFile(req.FilePath) + if err != nil { + return nil, err + } + + content := string(raw) + offset := req.Offset - 1 + if offset < 0 { + offset = 0 + } + limit := req.Limit + + if offset == 0 && limit <= 0 { + return &FileContent{Content: content}, nil + } + + start := 0 + for i := 0; i < offset; i++ { + idx := strings.IndexByte(content[start:], '\n') + if idx == -1 { + return &FileContent{Content: ""}, nil + } + start += idx + 1 + } + + if limit <= 0 { + return &FileContent{Content: content[start:]}, nil + } + + end := start + for i := 0; i < limit; i++ { + idx := strings.IndexByte(content[end:], '\n') + if idx == -1 { + return &FileContent{Content: content[start:]}, nil + } + end += idx + 1 + } + + return &FileContent{Content: strings.TrimSuffix(content[start:end], "\n")}, nil +} + +func (b *LocalBackend) GlobInfo(_ context.Context, req *GlobInfoRequest) ([]FileInfo, error) { + if req == nil || req.Pattern == "" || req.Path == "" { + return nil, fmt.Errorf("glob: invalid request") + } + + root := filepath.Clean(req.Path) + var matches []FileInfo + type item struct { + fi FileInfo + t time.Time + } + var tmp []item + + err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() { + return nil + } + + rel, err := filepath.Rel(root, path) + if err != nil { + return err + } + rel = filepath.ToSlash(rel) + + ok, err := doublestar.Match(req.Pattern, rel) + if err != nil { + return err + } + if !ok { + return nil + } + + st, err := os.Stat(path) + if err != nil { + return err + } + + tmp = append(tmp, item{ + fi: FileInfo{ + Path: path, + IsDir: false, + Size: st.Size(), + ModifiedAt: st.ModTime().Format(time.RFC3339Nano), + }, + t: st.ModTime(), + }) + return nil + }) + if err != nil { + return nil, err + } + + sort.Slice(tmp, func(i, j int) bool { return tmp[i].t.After(tmp[j].t) }) + matches = make([]FileInfo, 0, len(tmp)) + for _, it := range tmp { + matches = append(matches, it.fi) + } + return matches, nil +} + +func (b *LocalBackend) Write(_ context.Context, req *WriteRequest) error { + if req == nil || req.FilePath == "" { + return fmt.Errorf("write: invalid request") + } + path := filepath.Clean(req.FilePath) + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + + tmp := path + ".tmp" + if err := os.WriteFile(tmp, []byte(req.Content), 0o644); err != nil { + return err + } + return os.Rename(tmp, path) +} + +func (b *LocalBackend) Edit(ctx context.Context, req *EditRequest) error { + if req == nil || req.FilePath == "" { + return fmt.Errorf("edit: invalid request") + } + fc, err := b.Read(ctx, &ReadRequest{FilePath: req.FilePath}) + if err != nil { + return err + } + if req.OldString == "" { + return fmt.Errorf("edit: old string must be non-empty") + } + if req.OldString == req.NewString { + return fmt.Errorf("edit: new string must differ from old string") + } + + out := fc.Content + if req.ReplaceAll { + out = strings.ReplaceAll(out, req.OldString, req.NewString) + } else { + if strings.Count(out, req.OldString) != 1 { + return fmt.Errorf("edit: old string must appear exactly once when ReplaceAll is false") + } + out = strings.Replace(out, req.OldString, req.NewString, 1) + } + return b.Write(ctx, &WriteRequest{FilePath: req.FilePath, Content: out}) +} diff --git a/adk/middlewares/automemory/prompt.go b/adk/middlewares/automemory/prompt.go new file mode 100644 index 000000000..434fd9d95 --- /dev/null +++ b/adk/middlewares/automemory/prompt.go @@ -0,0 +1,588 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/cloudwego/eino/adk/internal" +) + +const ( + defaultMemoryInstructionWithIndex = `# Auto memory + +You have access to a persistent memory directory. Its contents persist across conversations. + +As you work, consult your memory files to build on previous experience. + +## How to save memories: +- Organize memory semantically by topic, not chronologically +- Use the Write and Edit tools to update your memory files +- When MEMORY.md is enabled, it is provided as a memory index reminder — content is truncated after configured line and byte limits, so keep it concise +- Create separate topic files (e.g., 'debugging.md', 'patterns.md') for detailed notes and link to them from MEMORY.md +- Update or remove memories that turn out to be wrong or outdated +- Do not write duplicate memories. First check if there is an existing memory you can update before writing a new one. + +## What to save: +- Stable patterns and conventions confirmed across multiple interactions +- Key architectural decisions, important file paths, and project structure +- User preferences for workflow, tools, and communication style +- Solutions to recurring problems and debugging insights + +## What NOT to save: +- Session-specific context (current task details, in-progress work, temporary state) +- Information that might be incomplete — verify against project docs before writing +- Anything that duplicates or contradicts existing AGENTS.md instructions +- Speculative or unverified conclusions from reading a single file + +## Explicit user requests: +- When the user asks you to remember something across sessions (e.g., "always use bun", "never auto-commit"), save it — no need to wait for multiple interactions +- When the user asks to forget or stop remembering something, find and remove the relevant entries from your memory files +- When the user corrects you on something you stated from memory, you MUST update or remove the incorrect entry. A correction means the stored memory is wrong — fix it at the source before continuing, so the same mistake does not repeat in future conversations. + +## Searching past context +- Search topic files inside the memory directory. Grep with pattern="" path="" glob="*.md" +- Use narrow search terms (error messages, file paths, function names) rather than broad keywords. + +` + + defaultAppendCurrentIndexTruncNotify = `WARNING: MEMORY.md was truncated (lines: {memory_lines}, limit: 200; byte limit: 4096). Move detailed content into separate topic files and keep MEMORY.md as a concise index.` + + defaultAppendEmptyIndexTemplate = `Your MEMORY.md is currently empty. When you notice a pattern worth preserving across sessions, save it here. Anything in MEMORY.md may be surfaced as a memory index reminder in future runs.` + + defaultTopicSelectionSystemPrompt = `You are selecting memories that will be useful to the agent as it processes a user's query. You will be given the user's query and a list of available memory files from the memory directory, with their displayed memory paths and descriptions. + +Return a list of memory paths exactly as shown in the available memories list, for the memories that will clearly be useful to the agent as it processes the user's query, up to the selection limit provided by the user message. Only include memories that you are certain will be helpful based on their path, name, description, or type. +- If you are unsure if a memory will be useful in processing the user's query, then do not include it in your list. Be selective and discerning. +- If there are no memories in the list that would clearly be useful, feel free to return an empty list. +- If a list of recently-used tools is provided, do not select memories that are usage reference or API documentation for those tools (the agent is already exercising them). DO still select memories containing warnings, gotchas, or known issues about those tools — active use is exactly when those matter.` + + defaultTopicSelectionUserPrompt = `Query: {user_query} + +Selection limit: {top_k} + +Available memories: +{available_memories} + +Recently used tools: +{tools}` + + defaultTopicMemoryTruncNotify = ` +> This memory file was truncated ({reason}). Use the Read tool to view the complete file at: {abs_path}` + + defaultMemoryInstructionChineseWithIndex = `# 自动记忆 + +你可以访问持久化的记忆目录。其中的内容会在不同会话之间保留。 + +在工作过程中，请查阅这些记忆文件，以便基于过去的经验继续推进。 + +## 如何保存记忆： +- 按主题组织记忆，而不是按时间顺序堆叠 +- 使用 Write 和 Edit 工具更新你的记忆文件 +- 启用 MEMORY.md 时，它会作为记忆索引 reminder 提供，其内容会按配置的行数和字节数限制截断，因此请保持简洁 +- 将详细内容写入单独的主题文件（例如 'debugging.md'、'patterns.md'），并在 MEMORY.md 中链接它们 +- 当某条记忆被证明错误或过时时，请更新或删除它 +- 不要写入重复记忆。创建新记忆前，先检查是否已有可更新的现有文件 + +## 应该保存什么： +- 已在多次交互中得到确认的稳定模式和约定 +- 关键架构决策、重要文件路径和项目结构 +- 用户在工作流、工具使用和沟通方式上的偏好 +- 可复用的问题解决经验与调试结论 + +## 不应保存什么： +- 仅属于当前会话的上下文（当前任务细节、进行中的工作、临时状态） +- 可能不完整的信息，在写入前应先根据项目文档核实 +- 与现有 AGENTS.md 指令重复或冲突的内容 +- 仅基于阅读单个文件得到的猜测性或未经验证的结论 + +## 用户的明确要求： +- 当用户明确要求你跨会话记住某件事时（例如“始终使用 bun”“不要自动提交”），应立即保存，无需等待多轮交互确认 +- 当用户要求你遗忘某件事或停止记忆时，找到对应条目并从记忆文件中删除 +- 当用户指出你基于记忆给出的内容有误时，你必须更新或删除错误条目。纠正意味着原有记忆已经错误，必须先从源头修正，避免今后重复犯错 + +## 如何检索历史上下文 +- 在记忆目录中搜索主题文件。使用 pattern="<搜索词>" path="<记忆目录路径>" glob="*.md" 进行 grep 搜索。 +- 尽量使用更窄的检索词，例如报错信息、文件路径、函数名，而不是宽泛关键词 + +` + + defaultAppendCurrentIndexTruncNotifyChinese = `警告：MEMORY.md 已被截断（总行数：{memory_lines}，限制：200 行；字节限制：4096）。请将详细内容迁移到独立的主题文件中，并让 MEMORY.md 只保留简洁索引。` + + defaultAppendEmptyIndexTemplateChinese = `你的 MEMORY.md 当前为空。当你发现值得跨会话保留的模式时，请把它写在这里。后续运行中，MEMORY.md 的内容可能会作为记忆索引 reminder 提供。` + + defaultTopicSelectionSystemPromptChinese = `你需要从记忆列表中选择对当前用户问题真正有帮助的记忆。你会拿到用户问题，以及来自记忆目录的可用记忆文件列表，列表中包含展示给你的记忆路径和描述。 + +请返回一个记忆路径列表，必须与可用记忆列表中展示的路径完全一致，列出那些在处理当前用户问题时显然有帮助的记忆文件，数量不能超过用户消息中给出的选择上限。只有在你能够基于路径、名称、描述或类型确认其确实有帮助时才选择。 +- 如果你不能确定某条记忆是否有帮助，就不要选它。请保持克制和甄别。 +- 如果列表中没有任何明显有帮助的记忆，可以返回空列表。 +- 如果提供了最近使用过的工具列表，不要选择那些仅包含这些工具使用说明或 API 文档的记忆（智能体已经在使用它们）。但如果记忆中包含这些工具的警告、坑点或已知问题，仍然应该选择，因为这些内容在实际调用时尤其重要。` + + defaultTopicSelectionUserPromptChinese = `问题：{user_query} + +选择上限：{top_k} + +可用记忆： +{available_memories} + +最近使用的工具： +{tools}` + + defaultTopicMemoryTruncNotifyChinese = ` +> 该记忆文件已被截断（{reason}）。请使用 Read 工具查看完整文件：{abs_path}` +) + +type memoryIndexPromptInfo struct { + FileName string + Path string + Content string + Empty bool + Truncated bool + Lines int + IncludeContent bool +} + +type memoryManifestPromptInfo struct { + Directory string + Files []memoryManifestFilePromptInfo +} + +type memoryManifestFilePromptInfo struct { + MemoryPath string + AbsPath string + Saved string + Description string +} + +func buildSystemMemoryInstruction(baseInstruction, memoryInstruction, memoryDirectory string) (string, error) { + return baseInstruction + "\n" + internal.SelectPrompt(internal.I18nPrompts{ + English: buildSystemMemoryInstructionEnglish(memoryInstruction, memoryDirectory), + Chinese: buildSystemMemoryInstructionChinese(memoryInstruction, memoryDirectory), + }), nil +} + +func buildSystemMemoryInstructionEnglish(memoryInstruction string, memoryDirectory string) string { + return strings.Join([]string{memoryInstruction, buildMemoryDirectoryManifestEnglish(memoryDirectory, nil)}, "\n") +} + +func buildSystemMemoryInstructionChinese(memoryInstruction string, memoryDirectory string) string { + return strings.Join([]string{memoryInstruction, buildMemoryDirectoryManifestChinese(memoryDirectory, nil)}, "\n") +} + +func buildMemoryDirectoryManifestEnglish(memoryDirectory string, index *memoryIndexPromptInfo) string { + lines := []string{ + "## Memory directory", + "", + fmt.Sprintf("Path: %s", memoryDirectory), + } + if index != nil { + lines = append(lines, fmt.Sprintf("Index file path: %s", index.Path), "") + if block := buildMemoryIndexBlockEnglish(*index); block != "" { + lines = append(lines, block) + } + } + return strings.Join(lines, "\n") +} + +func buildMemoryDirectoryManifestChinese(memoryDirectory string, index *memoryIndexPromptInfo) string { + lines := []string{ + "## 记忆目录", + "", + fmt.Sprintf("路径：%s", memoryDirectory), + } + if index != nil { + lines = append(lines, fmt.Sprintf("索引文件路径：%s", index.Path), "") + if block := buildMemoryIndexBlockChinese(*index); block != "" { + lines = append(lines, block) + } + } + return strings.Join(lines, "\n") +} + +func buildMemoryIndexBlockEnglish(index memoryIndexPromptInfo) string { + if !index.IncludeContent { + return "" + } + lines := []string{fmt.Sprintf("#### Index file content: %s", index.FileName)} + if index.Empty { + lines = append(lines, getAppendEmptyIndexTemplate()) + } else { + lines = append(lines, index.Content) + if index.Truncated { + lines = append(lines, strings.ReplaceAll(getAppendCurrentIndexTruncNotify(), "{memory_lines}", fmt.Sprintf("%d", index.Lines))) + } + } + return strings.Join(lines, "\n") +} + +func buildMemoryIndexBlockChinese(index memoryIndexPromptInfo) string { + if !index.IncludeContent { + return "" + } + lines := []string{fmt.Sprintf("#### 索引文件内容：%s", index.FileName)} + if index.Empty { + lines = append(lines, getAppendEmptyIndexTemplate()) + } else { + lines = append(lines, index.Content) + if index.Truncated { + lines = append(lines, strings.ReplaceAll(getAppendCurrentIndexTruncNotify(), "{memory_lines}", fmt.Sprintf("%d", index.Lines))) + } + } + return strings.Join(lines, "\n") +} + +func buildExtractAutoOnlyPrompt(memoryStores string, newMessageCount int, existingMemories string, savePolicyInstruction string) string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: buildExtractAutoOnlyPromptEnglish(memoryStores, newMessageCount, existingMemories, savePolicyInstruction), + Chinese: buildExtractAutoOnlyPromptChinese(memoryStores, newMessageCount, existingMemories, savePolicyInstruction), + }) +} + +func buildMemoryDirectoryManifest(memoryDirectory string, index *memoryIndexPromptInfo) string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: buildMemoryDirectoryManifestEnglish(memoryDirectory, index), + Chinese: buildMemoryDirectoryManifestChinese(memoryDirectory, index), + }) +} + +func buildMemoryIndexReminder(index memoryIndexPromptInfo) string { + return "\n" + internal.SelectPrompt(internal.I18nPrompts{ + English: buildMemoryIndexReminderEnglish(index), + Chinese: buildMemoryIndexReminderChinese(index), + }) +} + +func buildTopicMemoryReminder(topics []topicMemoryPromptInfo) string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: buildTopicMemoryReminderEnglish(topics), + Chinese: buildTopicMemoryReminderChinese(topics), + }) +} + +func buildTopicMemoryReminderEnglish(topics []topicMemoryPromptInfo) string { + lines := []string{ + "", + } + for i, topic := range topics { + lines = append(lines, + fmt.Sprintf("", i+1), + fmt.Sprintf("Contents of %s (saved %s):", filepath.Join(topic.MemoryDirectory, topic.Path), topic.Saved), + topic.Content, + fmt.Sprintf("", i+1), + "", + ) + } + lines = append(lines, "") + return strings.Join(lines, "\n") +} + +func buildTopicMemoryReminderChinese(topics []topicMemoryPromptInfo) string { + lines := []string{ + "", + "主题记忆是本次查询相关的长期记忆文件。请将它们作为当前轮次的辅助上下文使用，其中可能包含稳定的用户偏好、项目约定或此前保存的事实；不要用它们替代当前用户请求。", + "", + } + for i, topic := range topics { + lines = append(lines, + fmt.Sprintf("", i+1), + fmt.Sprintf("主题记忆 %s 内容 (更新于 %s): ", filepath.Join(topic.MemoryDirectory, topic.Path), topic.Saved), + topic.Content, + fmt.Sprintf("", i+1), + "", + ) + } + lines = append(lines, "") + return strings.Join(lines, "\n") +} + +func buildMemoryIndexReminderEnglish(index memoryIndexPromptInfo) string { + return strings.Join([]string{ + "", + "As you answer the user's questions, you can use the following context:", + "# Memory Index", + renderMemoryIndexContentEnglish(index), + "", + "IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.", + "", + }, "\n") +} + +func buildMemoryIndexReminderChinese(index memoryIndexPromptInfo) string { + return strings.Join([]string{ + "", + "在回答用户问题时，您可以使用以下上下:", + "# 记忆索引文件", + renderMemoryIndexContentChinese(index), + "", + "重要提示: 此上下文未必与您的任务相关。除非与任务高度相关，否则不应该对此上下文作出回应。", + "", + }, "\n") +} + +func renderMemoryIndexContentEnglish(index memoryIndexPromptInfo) string { + if index.Empty { + return fmt.Sprintf("Contents of %s (user's auto-memory, persists across conversations) is currently empty.", index.Path) + } + + lines := []string{ + fmt.Sprintf("Contents of %s (user's auto-memory, persists across conversations):", index.Path), + "", + index.Content, + } + if index.Truncated { + lines = append(lines, strings.ReplaceAll(getAppendCurrentIndexTruncNotify(), "{memory_lines}", fmt.Sprintf("%d", index.Lines))) + } + + return strings.Join(lines, "\n") +} + +func renderMemoryIndexContentChinese(index memoryIndexPromptInfo) string { + if index.Empty { + return fmt.Sprintf("文件 %s（用户的自动记忆内容，在会话中持续存在）内容为空。", index.Path) + } + + lines := []string{ + fmt.Sprintf("文件 %s（用户的自动记忆内容，在会话中持续存在）内容:", index.Path), + "", + index.Content, + } + if index.Truncated { + lines = append(lines, strings.ReplaceAll(getAppendCurrentIndexTruncNotify(), "{memory_lines}", fmt.Sprintf("%d", index.Lines))) + } + + return strings.Join(lines, "\n") +} + +func buildExtractionMemoryManifest(manifest memoryManifestPromptInfo) string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: buildExtractionMemoryManifestEnglish(manifest), + Chinese: buildExtractionMemoryManifestChinese(manifest), + }) +} + +func buildExtractionMemoryManifestEnglish(manifest memoryManifestPromptInfo) string { + lines := []string{fmt.Sprintf("Memory directory: %s", manifest.Directory)} + if len(manifest.Files) == 0 { + lines = append(lines, "- No existing memory files.") + return strings.Join(lines, "\n") + } + for _, file := range manifest.Files { + if file.Description != "" { + lines = append(lines, fmt.Sprintf("- %s (path: %s, saved %s): %s", file.MemoryPath, file.AbsPath, file.Saved, file.Description)) + } else { + lines = append(lines, fmt.Sprintf("- %s (path: %s, saved %s)", file.MemoryPath, file.AbsPath, file.Saved)) + } + } + return strings.Join(lines, "\n") +} + +func buildExtractionMemoryManifestChinese(manifest memoryManifestPromptInfo) string { + lines := []string{fmt.Sprintf("记忆目录：%s", manifest.Directory)} + if len(manifest.Files) == 0 { + lines = append(lines, "- 暂无已有 memory 文件。") + return strings.Join(lines, "\n") + } + for _, file := range manifest.Files { + if file.Description != "" { + lines = append(lines, fmt.Sprintf("- %s（路径：%s，保存时间：%s）：%s", file.MemoryPath, file.AbsPath, file.Saved, file.Description)) + } else { + lines = append(lines, fmt.Sprintf("- %s（路径：%s，保存时间：%s）", file.MemoryPath, file.AbsPath, file.Saved)) + } + } + return strings.Join(lines, "\n") +} + +func joinLines(lines []string) string { + if len(lines) == 0 { + return "" + } + var b strings.Builder + b.WriteString(lines[0]) + for i := 1; i < len(lines); i++ { + b.WriteString("\n") + b.WriteString(lines[i]) + } + return b.String() +} + +func getDefaultMemoryInstruction() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultMemoryInstructionWithIndex, + Chinese: defaultMemoryInstructionChineseWithIndex, + }) +} + +func getAppendCurrentIndexTruncNotify() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultAppendCurrentIndexTruncNotify, + Chinese: defaultAppendCurrentIndexTruncNotifyChinese, + }) +} + +func getAppendEmptyIndexTemplate() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultAppendEmptyIndexTemplate, + Chinese: defaultAppendEmptyIndexTemplateChinese, + }) +} + +func getTopicSelectionSystemPrompt() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultTopicSelectionSystemPrompt, + Chinese: defaultTopicSelectionSystemPromptChinese, + }) +} + +func getTopicSelectionUserPrompt() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultTopicSelectionUserPrompt, + Chinese: defaultTopicSelectionUserPromptChinese, + }) +} + +func getTopicMemoryTruncNotify() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultTopicMemoryTruncNotify, + Chinese: defaultTopicMemoryTruncNotifyChinese, + }) +} + +func buildExtractHowToSaveEnglish() []string { + return []string{ + "## How to save memories", + "", + "Saving a memory is a two-step process:", + "", + "Step 1 — write the memory to its own file.", + "Step 2 — add a pointer to that file in MEMORY.md. MEMORY.md is an index, not the memory body.", + "", + "- Keep MEMORY.md concise because it is surfaced as a memory index reminder.", + "- Organize memory semantically by topic, not chronologically.", + "- Update or remove memories that turn out to be wrong or outdated.", + "- Do not write duplicate memories.", + } +} + +func buildExtractHowToSaveChinese() []string { + return []string{ + "## 如何保存记忆", + "", + "保存记忆分为两步：", + "", + "第 1 步：将记忆写入独立文件。", + "第 2 步：在 MEMORY.md 中添加指向该文件的索引。MEMORY.md 只是索引，不应存放记忆正文。", + "", + "- 保持 MEMORY.md 简洁，因为它会作为记忆索引 reminder 提供。", + "- 按主题组织记忆，而不是按时间顺序堆叠。", + "- 当记忆被证明错误或过时时，要及时更新或删除。", + "- 不要写入重复记忆。", + } +} + +func buildExtractSavePolicyEnglish(custom string) []string { + if strings.TrimSpace(custom) != "" { + return strings.Split(strings.TrimSpace(custom), "\n") + } + return []string{ + "## What to save", + "- Stable patterns and conventions confirmed across multiple interactions", + "- Important file paths, architectural decisions, and user preferences", + "- Recurring debugging insights and known gotchas", + "", + "## What NOT to save", + "- Session-specific temporary state or current task details", + "- Secrets, credentials, or personal data", + "- Speculative or unverified conclusions", + } +} + +func buildExtractSavePolicyChinese(custom string) []string { + if strings.TrimSpace(custom) != "" { + return strings.Split(strings.TrimSpace(custom), "\n") + } + return []string{ + "## 应该保存什么", + "- 已在多次交互中得到确认的稳定模式和约定", + "- 重要文件路径、架构决策和用户偏好", + "- 可复用的调试经验与已知坑点", + "", + "## 不应保存什么", + "- 仅属于当前会话的临时状态或当前任务细节", + "- 密钥、凭据或个人数据", + "- 猜测性或未经验证的结论", + } +} + +func buildExtractAutoOnlyPromptEnglish(memoryStores string, newMessageCount int, existingMemories string, savePolicyInstruction string) string { + manifest := "" + if existingMemories != "" { + manifest = fmt.Sprintf("\n\n## Existing memory files\n\n%s\n\nCheck this list before writing — update an existing file rather than creating a duplicate.", existingMemories) + } + + howToSave := buildExtractHowToSaveEnglish() + savePolicy := buildExtractSavePolicyEnglish(savePolicyInstruction) + + parts := []string{ + fmt.Sprintf("You are now acting as the memory extraction subagent. Analyze only the most recent ~%d messages above and use them to update persistent memory.", newMessageCount), + "", + memoryStores, + "", + "Available tools: read_file, glob, write_file, edit_file. Only paths inside the memory directory are allowed. Use absolute paths or paths relative to the memory directory when reading or writing memory files. All other tools are denied.", + "", + "You have a limited turn budget. read_file should happen first for every file you may update, then write_file/edit_file should happen after that. Do not interleave read and write across many turns.", + "", + fmt.Sprintf("You MUST only use content from the last ~%d messages to update memories. Do not investigate code or verify against source files further.", newMessageCount) + manifest, + "", + "If the user explicitly asks you to remember something, save it immediately. If they ask you to forget something, find and remove the relevant memory.", + "", + } + parts = append(parts, savePolicy...) + parts = append(parts, "") + parts = append(parts, howToSave...) + return joinLines(parts) +} + +func buildExtractAutoOnlyPromptChinese(memoryStores string, newMessageCount int, existingMemories string, savePolicyInstruction string) string { + manifest := "" + if existingMemories != "" { + manifest = fmt.Sprintf("\n\n## 现有记忆文件\n\n%s\n\n写入前请先检查这份列表，优先更新已有文件，而不是创建重复记忆。", existingMemories) + } + + howToSave := buildExtractHowToSaveChinese() + savePolicy := buildExtractSavePolicyChinese(savePolicyInstruction) + + parts := []string{ + fmt.Sprintf("你现在扮演记忆提取子智能体。只分析上方最近约 %d 条消息，并用它们来更新持久化记忆。", newMessageCount), + "", + memoryStores, + "", + "可用工具：read_file、glob、write_file、edit_file。只允许访问记忆目录内的路径。读写记忆文件时请使用绝对路径，或使用相对记忆目录的路径。其他工具均禁止使用。", + "", + "你的轮次预算有限。对于每个可能更新的文件，应先 read_file，再进行 write_file/edit_file；不要在多轮里交叉读写大量文件。", + "", + fmt.Sprintf("你必须只使用最近约 %d 条消息中的内容来更新记忆。不要继续调查代码，也不要再去源码中额外验证。", newMessageCount) + manifest, + "", + "如果用户明确要求你记住某件事，请立即保存；如果用户要求遗忘某件事，请找到对应记忆并删除。", + "", + } + parts = append(parts, savePolicy...) + parts = append(parts, "") + parts = append(parts, howToSave...) + return joinLines(parts) +} diff --git a/adk/middlewares/automemory/utils.go b/adk/middlewares/automemory/utils.go new file mode 100644 index 000000000..e8c7d6bf1 --- /dev/null +++ b/adk/middlewares/automemory/utils.go @@ -0,0 +1,875 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package automemory + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "sort" + "strings" + "time" + + "gopkg.in/yaml.v3" + + "github.com/cloudwego/eino/adk" + adkfs "github.com/cloudwego/eino/adk/middlewares/filesystem" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +func applyReadDefaults[M adk.MessageType](cfg *Config[M]) { + if cfg.Read.Mode == "" { + cfg.Read.Mode = ReadModeSync + } + if cfg.Read.Index == nil { + cfg.Read.Index = &IndexConfig{} + } + if cfg.Read.Index.FileName == "" { + cfg.Read.Index.FileName = memoryIndexFileName + } + if cfg.Read.Index.MaxLines <= 0 { + cfg.Read.Index.MaxLines = defaultIndexMaxLines + } + if cfg.Read.Index.MaxBytes <= 0 { + cfg.Read.Index.MaxBytes = defaultIndexMaxBytes + } + if cfg.Read.Model == nil { + cfg.Read.Model = cfg.Model + } + if cfg.Read.TopicSelection == nil { + cfg.Read.TopicSelection = &TopicSelectionConfig{} + } + if cfg.Read.TopicSelection.Enable == nil { + cfg.Read.TopicSelection.Enable = boolPtr(true) + } + if cfg.Read.TopicSelection.TopK <= 0 { + cfg.Read.TopicSelection.TopK = defaultTopicTopK + } + if cfg.Read.TopicSelection.CandidateGlob == "" { + cfg.Read.TopicSelection.CandidateGlob = CandidateGlobPattern + } + if cfg.Read.TopicSelection.CandidateLimit <= 0 { + cfg.Read.TopicSelection.CandidateLimit = defaultCandidateLimit + } + if cfg.Read.TopicSelection.CandidatePreviewLines <= 0 { + cfg.Read.TopicSelection.CandidatePreviewLines = defaultCandidatePreviewLine + } + if cfg.Read.TopicSelection.MaxLines <= 0 { + cfg.Read.TopicSelection.MaxLines = defaultTopicMaxLines + } + if cfg.Read.TopicSelection.MaxBytes <= 0 { + cfg.Read.TopicSelection.MaxBytes = defaultTopicMaxBytes + } + if cfg.Read.TopicSelection.MaxTotalBytes <= 0 { + cfg.Read.TopicSelection.MaxTotalBytes = defaultTopicMaxTotalBytes + } + + if cfg.Write == nil { + cfg.Write = &WriteConfig[M]{Mode: WriteModeDisabled} + } + if cfg.Write.Mode == "" { + cfg.Write.Mode = WriteModeDisabled + } + if cfg.Write.Model == nil { + cfg.Write.Model = cfg.Model + } + if cfg.Write.MaxTurns <= 0 { + cfg.Write.MaxTurns = defaultMemoryWriteMaxTurns + } + + if cfg.Coordination == nil { + cfg.Coordination = &CoordinationConfig[M]{} + } + if cfg.Coordination.Coordinator == nil { + cfg.Coordination.Coordinator = NewLocalCoordinator() + } + if cfg.Coordination.LockTTL <= 0 { + cfg.Coordination.LockTTL = 2 * time.Minute + } +} + +func cloneConfig[M adk.MessageType](cfg *Config[M]) *Config[M] { + if cfg == nil { + return nil + } + + cp := *cfg + if cfg.Read != nil { + readCopy := *cfg.Read + cp.Read = &readCopy + if cfg.Read.Index != nil { + indexCopy := *cfg.Read.Index + cp.Read.Index = &indexCopy + } + if cfg.Read.TopicSelection != nil { + topicSelectionCopy := *cfg.Read.TopicSelection + cp.Read.TopicSelection = &topicSelectionCopy + } + } + if cfg.Write != nil { + writeCopy := *cfg.Write + cp.Write = &writeCopy + } + if cfg.Coordination != nil { + coordinationCopy := *cfg.Coordination + cp.Coordination = &coordinationCopy + } + return &cp +} + +func linesOrSizeTrunc(content string, lines, size int) (newContent string, reason string, truncated bool) { + linesTrunc := func(content string, lines int) { + sp := strings.Split(content, "\n") + if len(sp) > lines { + newContent = strings.Join(sp[:lines], "\n") + reason = fmt.Sprintf("first %d lines", lines) + truncated = true + } else { + newContent = content + } + } + + sizeTrunc := func(content string, size int) { + if len(content) > size { + newContent = content[:size] + reason = fmt.Sprintf("%d byte limit", size) + truncated = true + } else { + newContent = content + } + } + + if lines == 0 && size == 0 { + return content, "", false + } else if lines == 0 { + sizeTrunc(content, size) + } else if size == 0 { + linesTrunc(content, lines) + } else { + linesTrunc(content, lines) + sizeTrunc(newContent, size) + } + return +} + +func isFileNotFoundContent(content string) bool { + return strings.HasPrefix(strings.TrimSpace(content), "File not found: ") +} + +func boolPtr(v bool) *bool { + return &v +} + +func parseFrontmatter(md string) (fm topicFrontmatter, ok bool) { + s := strings.TrimLeft(md, "\ufeff \t\r\n") + if !strings.HasPrefix(s, "---\n") && !strings.HasPrefix(s, "---\r\n") { + return topicFrontmatter{}, false + } + parts := strings.SplitN(s, "\n---", 2) + if len(parts) != 2 { + return topicFrontmatter{}, false + } + yml := strings.TrimPrefix(parts[0], "---\n") + if err := yaml.Unmarshal([]byte(yml), &fm); err != nil { + return topicFrontmatter{}, false + } + return fm, true +} + +func describeTopicCandidate(content string) string { + desc := "" + if fm, ok := parseFrontmatter(content); ok { + switch { + case strings.TrimSpace(fm.Description) != "": + desc = strings.TrimSpace(fm.Description) + case strings.TrimSpace(fm.Name) != "": + desc = strings.TrimSpace(fm.Name) + } + if strings.TrimSpace(fm.Type) != "" { + if desc == "" { + desc = "type=" + strings.TrimSpace(fm.Type) + } else { + desc = desc + " (type=" + strings.TrimSpace(fm.Type) + ")" + } + } + } + if desc == "" { + snippet, _, _ := linesOrSizeTrunc(content, 3, 256) + desc = strings.TrimSpace(snippet) + } + return desc +} + +func collectToolNames[M adk.MessageType](msgs []M) []string { + dedupTools := make(map[string]struct{}) + for _, msg := range msgs { + for _, name := range messageToolNames(msg) { + dedupTools[name] = struct{}{} + } + } + tools := make([]string, 0, len(dedupTools)) + for t := range dedupTools { + tools = append(tools, t) + } + sort.Strings(tools) + return tools +} + +func topicSelectionToolInfo() *schema.ToolInfo { + return &schema.ToolInfo{ + Name: topicSelectionToolName, + Desc: "Select which memory files to surface for the current query. Return selected_memories as memory paths exactly as shown in the available memories list.", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "selected_memories": { + Type: schema.Array, + Desc: "Memory paths exactly as shown in the available memories list, e.g. \"user_profile/preferences.md\" or \"project_context/notes/patterns.md\".", + Required: true, + ElemInfo: &schema.ParameterInfo{Type: schema.String}, + }, + }), + } +} + +func parseTopicSelectionFromToolCall[M adk.MessageType](msg M, valid map[string]struct{}) ([]string, error) { + toolCalls := messageToolCalls(msg) + if len(toolCalls) == 0 { + return nil, fmt.Errorf("no tool calls") + } + tc := toolCalls[0] + if tc.Function.Name != topicSelectionToolName { + return nil, fmt.Errorf("unexpected tool call: %s", tc.Function.Name) + } + var parsed topicSelectionResp + if err := json.Unmarshal([]byte(tc.Function.Arguments), &parsed); err != nil { + return nil, err + } + out := normalizeSelected(parsed.SelectedMemories) + filtered := make([]string, 0, len(out)) + for _, p := range out { + if _, ok := valid[p]; ok { + filtered = append(filtered, p) + } + } + return filtered, nil +} + +func normalizeSelected(in []string) []string { + out := make([]string, 0, len(in)) + seen := make(map[string]struct{}, len(in)) + for _, s := range in { + s = strings.TrimSpace(s) + s = strings.TrimPrefix(s, "./") + s = filepath.ToSlash(s) + if s == "" { + continue + } + if _, ok := seen[s]; ok { + continue + } + seen[s] = struct{}{} + out = append(out, s) + } + return out +} + +func isNilMessage[M adk.MessageType](msg M) bool { + var zero M + return any(msg) == any(zero) +} + +func isUserRole[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m != nil && m.Role == schema.User + case *schema.AgenticMessage: + return m != nil && m.Role == schema.AgenticRoleTypeUser + default: + panic("unreachable") + } +} + +func isAssistantRole[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m != nil && m.Role == schema.Assistant + case *schema.AgenticMessage: + return m != nil && m.Role == schema.AgenticRoleTypeAssistant + default: + panic("unreachable") + } +} + +func userMessageTextContent[M adk.MessageType](msg M) string { + switch m := any(msg).(type) { + case *schema.Message: + if m == nil { + return "" + } + if len(m.UserInputMultiContent) == 0 { + return m.Content + } + parts := make([]string, 0, len(m.UserInputMultiContent)) + for _, part := range m.UserInputMultiContent { + if part.Type == schema.ChatMessagePartTypeText && part.Text != "" { + parts = append(parts, part.Text) + } + } + if len(parts) > 0 { + return strings.Join(parts, "\n") + } + return m.Content + case *schema.AgenticMessage: + if m == nil { + return "" + } + parts := make([]string, 0, len(m.ContentBlocks)) + for _, block := range m.ContentBlocks { + if block != nil && block.UserInputText != nil { + parts = append(parts, block.UserInputText.Text) + } + } + return strings.Join(parts, "\n") + default: + panic("unreachable") + } +} + +func getMsgExtra[M adk.MessageType](msg M) map[string]any { + switch m := any(msg).(type) { + case *schema.Message: + if m == nil { + return nil + } + return m.Extra + case *schema.AgenticMessage: + if m == nil { + return nil + } + return m.Extra + default: + panic("unreachable") + } +} + +func copyAndSetMsgExtra[M adk.MessageType](msg M, key string, value any) { + existing := getMsgExtra(msg) + newExtra := make(map[string]any, len(existing)+1) + for k, v := range existing { + newExtra[k] = v + } + newExtra[key] = value + + switch m := any(msg).(type) { + case *schema.Message: + m.Extra = newExtra + case *schema.AgenticMessage: + m.Extra = newExtra + default: + panic("unreachable") + } +} + +func makeUserMsg[M adk.MessageType](text string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(text)).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(text)).(M) + default: + panic("unreachable") + } +} + +func makeSystemMsg[M adk.MessageType](text string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.SystemMessage(text)).(M) + case *schema.AgenticMessage: + return any(schema.SystemAgenticMessage(text)).(M) + default: + panic("unreachable") + } +} + +func makeToolChoiceForced[M adk.MessageType](name string) model.Option { + var zero M + switch any(zero).(type) { + case *schema.Message: + return model.WithToolChoice(schema.ToolChoiceForced, name) + case *schema.AgenticMessage: + return model.WithAgenticToolChoice(&schema.AgenticToolChoice{ + Type: schema.ToolChoiceForced, + Forced: &schema.AgenticForcedToolChoice{ + Tools: []*schema.AllowedTool{{FunctionName: name}}, + }, + }) + default: + panic("unreachable") + } +} + +func messageToolCalls[M adk.MessageType](msg M) []schema.ToolCall { + switch m := any(msg).(type) { + case *schema.Message: + if m == nil { + return nil + } + return m.ToolCalls + case *schema.AgenticMessage: + if m == nil { + return nil + } + out := make([]schema.ToolCall, 0, len(m.ContentBlocks)) + for _, block := range m.ContentBlocks { + if block == nil || block.FunctionToolCall == nil { + continue + } + out = append(out, schema.ToolCall{ + ID: block.FunctionToolCall.CallID, + Type: "function", + Function: schema.FunctionCall{ + Name: block.FunctionToolCall.Name, + Arguments: block.FunctionToolCall.Arguments, + }, + }) + } + return out + default: + panic("unreachable") + } +} + +func messageToolNames[M adk.MessageType](msg M) []string { + switch m := any(msg).(type) { + case *schema.Message: + if m == nil || m.Role != schema.Tool || m.ToolName == "" { + return nil + } + return []string{m.ToolName} + case *schema.AgenticMessage: + if m == nil { + return nil + } + var out []string + for _, block := range m.ContentBlocks { + if block == nil || block.FunctionToolResult == nil || block.FunctionToolResult.Name == "" { + continue + } + out = append(out, block.FunctionToolResult.Name) + } + return out + default: + panic("unreachable") + } +} + +func hasTopicMemoryInjected[M adk.MessageType](msgs []M) bool { + for _, msg := range msgs { + if isTopicMemoryMessage(msg) { + return true + } + } + return false +} + +func hasMemoryIndexInjected[M adk.MessageType](msgs []M) bool { + for _, msg := range msgs { + if isMemoryIndexMessage(msg) { + return true + } + } + return false +} + +func insertMessagesBeforeLastUserQuery[M adk.MessageType](msgs []M, inserts []M) []M { + if len(inserts) == 0 { + return msgs + } + idx := lastUserQueryMessageIndex(msgs) + if idx < 0 { + idx = len(msgs) + } + out := make([]M, 0, len(msgs)+len(inserts)) + out = append(out, msgs[:idx]...) + out = append(out, inserts...) + out = append(out, msgs[idx:]...) + return out +} + +func lastUserQueryMessageIndex[M adk.MessageType](msgs []M) int { + for i := len(msgs) - 1; i >= 0; i-- { + msg := msgs[i] + if isNilMessage(msg) || !isUserRole(msg) || isAutomemoryReminderMessage(msg) { + continue + } + return i + } + return -1 +} + +func isAutomemoryReminderMessage[M adk.MessageType](m M) bool { + if isTopicMemoryMessage(m) || isMemoryIndexMessage(m) { + return true + } + if isNilMessage(m) || !isUserRole(m) { + return false + } + return strings.HasPrefix(strings.TrimSpace(userMessageTextContent(m)), "") +} + +func isTopicMemoryMessage[M adk.MessageType](m M) bool { + if isNilMessage(m) || !isUserRole(m) { + return false + } + if extra := getMsgExtra(m); extra != nil { + if v, ok := extra[memoryExtraKey]; ok { + if isTopicMemoryExtra(v) { + return true + } + } + } + content := userMessageTextContent(m) + return strings.Contains(content, "") && !strings.Contains(content, "") +} + +func isMemoryIndexMessage[M adk.MessageType](m M) bool { + if isNilMessage(m) || !isUserRole(m) { + return false + } + if extra := getMsgExtra(m); extra != nil { + if v, ok := extra[memoryExtraKey]; ok { + if isMemoryIndexExtra(v) { + return true + } + } + } + return strings.Contains(userMessageTextContent(m), "") +} + +func isTopicMemoryExtra(v any) bool { + switch meta := v.(type) { + case *memoryExtra: + return meta != nil && (meta.Type == "memory" || meta.Type == "topic_memory") + case map[string]any: + typ, _ := meta["type"].(string) + return typ == "memory" || typ == "topic_memory" + default: + return false + } +} + +func isMemoryIndexExtra(v any) bool { + switch meta := v.(type) { + case *memoryExtra: + return meta != nil && meta.Type == "memory_index" + case map[string]any: + typ, _ := meta["type"].(string) + return typ == "memory_index" + default: + return false + } +} + +func newMemoryMessage[M adk.MessageType](content string) M { + msg := makeUserMsg[M](content) + copyAndSetMsgExtra(msg, memoryExtraKey, &memoryExtra{Type: "memory"}) + return msg +} + +func newMemoryIndexMessage[M adk.MessageType](content string) M { + msg := makeUserMsg[M](content) + copyAndSetMsgExtra(msg, memoryExtraKey, &memoryExtra{Type: "memory_index"}) + return msg +} + +func ensureMemoryMsgUnchanged[M adk.MessageType](state *adk.TypedChatModelAgentState[M], expectedContent string) *adk.TypedChatModelAgentState[M] { + if state == nil || strings.TrimSpace(expectedContent) == "" { + return state + } + changed := false + out := *state + out.Messages = append([]M{}, state.Messages...) + + for i, m := range out.Messages { + if !isTopicMemoryMessage(m) { + continue + } + extra := getMsgExtra(m) + if userMessageTextContent(m) != expectedContent || extra == nil || extra[memoryExtraKey] == nil { + out.Messages[i] = newMemoryMessage[M](expectedContent) + changed = true + } + } + if !changed { + return state + } + return &out +} + +func extractFilePath(args string) (string, bool) { + var m map[string]any + if err := json.Unmarshal([]byte(args), &m); err != nil { + return "", false + } + if v, ok := m["file_path"]; ok { + if s, ok := v.(string); ok && s != "" { + return s, true + } + } + if v, ok := m["filePath"]; ok { + if s, ok := v.(string); ok && s != "" { + return s, true + } + } + return "", false +} + +func isPathWithinMemoryDir(memDir string, filePath string) bool { + if memDir == "" || filePath == "" { + return false + } + md := filepath.Clean(memDir) + fp := filepath.Clean(filePath) + if !filepath.IsAbs(fp) { + fp = filepath.Join(md, fp) + fp = filepath.Clean(fp) + } + if fp == md { + return true + } + sep := string(filepath.Separator) + return strings.HasPrefix(fp, md+sep) +} + +func getWriteCursorFromMessages[M adk.MessageType](msgs []M) int { + for i := len(msgs) - 1; i >= 0; i-- { + m := msgs[i] + extra := getMsgExtra(m) + if isNilMessage(m) || extra == nil { + continue + } + v, ok := extra[memoryExtraKey] + if !ok { + continue + } + switch meta := v.(type) { + case *memoryExtra: + if meta != nil && meta.Type == "write_cursor" { + return meta.Cursor + } + case map[string]any: + if typ, _ := meta["type"].(string); typ != "write_cursor" { + continue + } + switch c := meta["cursor"].(type) { + case int: + return c + case int64: + return int(c) + case float64: + return int(c) + } + } + } + return 0 +} + +func markWriteCursor[M adk.MessageType](state *adk.TypedChatModelAgentState[M], cursor int) *adk.TypedChatModelAgentState[M] { + if state == nil || len(state.Messages) == 0 { + return state + } + last := state.Messages[len(state.Messages)-1] + if isNilMessage(last) { + return state + } + + copyAndSetMsgExtra(last, memoryExtraKey, &memoryExtra{ + Type: "write_cursor", + Cursor: cursor, + }) + + return state +} + +func countModelVisibleMessages[M adk.MessageType](msgs []M) int { + n := 0 + for _, m := range msgs { + if isNilMessage(m) { + continue + } + if isUserRole(m) || isAssistantRole(m) { + n++ + } + } + return n +} + +func buildPendingSnapshot[M adk.MessageType](messages []M, cursor int, toolInfos []*schema.ToolInfo) (*PendingSnapshot, error) { + raw, err := json.Marshal(messages) + if err != nil { + return nil, err + } + var rawToolInfos json.RawMessage + if toolInfos != nil { + rawToolInfos, err = json.Marshal(toolInfos) + if err != nil { + return nil, err + } + } + return &PendingSnapshot{Cursor: cursor, Messages: raw, ToolInfos: rawToolInfos}, nil +} + +func decodePendingSnapshot[M adk.MessageType](snapshot *PendingSnapshot) ([]M, int, []*schema.ToolInfo, error) { + if snapshot == nil { + return nil, 0, nil, nil + } + var msgs []M + if err := json.Unmarshal(snapshot.Messages, &msgs); err != nil { + return nil, 0, nil, err + } + var toolInfos []*schema.ToolInfo + if len(snapshot.ToolInfos) > 0 { + if err := json.Unmarshal(snapshot.ToolInfos, &toolInfos); err != nil { + return nil, 0, nil, err + } + } + return msgs, snapshot.Cursor, toolInfos, nil +} + +func hasMemoryWritesSince[M adk.MessageType](msgs []M, cursor int, memoryDirectory string) bool { + if cursor < 0 { + cursor = 0 + } + for _, msg := range msgs[cursor:] { + if isNilMessage(msg) || !isAssistantRole(msg) { + continue + } + for _, tc := range messageToolCalls(msg) { + if tc.Function.Name != adkfs.ToolNameWriteFile && tc.Function.Name != adkfs.ToolNameEditFile { + continue + } + if fp, ok := extractFilePath(tc.Function.Arguments); ok && isPathWithinMemoryDir(memoryDirectory, fp) { + return true + } + } + } + return false +} + +func countModelVisibleMessagesSince[M adk.MessageType](msgs []M, cursor int) int { + if cursor < 0 { + cursor = 0 + } + if cursor >= len(msgs) { + return 0 + } + return countModelVisibleMessages(msgs[cursor:]) +} + +func parseRFC3339NanoBestEffort(s string) time.Time { + if s == "" { + return time.Time{} + } + if t, err := time.Parse(time.RFC3339Nano, s); err == nil { + return t + } + if t, err := time.Parse(time.RFC3339, s); err == nil { + return t + } + return time.Time{} +} + +func (m *middleware[M]) coordinatorKey(sessionID string) string { + if sessionID == "" || m == nil || m.resolvedMemoryDirectory == "" { + return "" + } + return m.resolvedMemoryDirectory + "::" + sessionID +} + +func (m *middleware[M]) topicSelectionEnabled() bool { + return m != nil && m.cfg != nil && m.cfg.Read != nil && + topicSelectionConfigEnabled(m.cfg.Read.TopicSelection) && m.topicSelectionModel != nil +} + +func topicSelectionConfigEnabled(cfg *TopicSelectionConfig) bool { + return cfg != nil && cfg.Enable != nil && *cfg.Enable +} + +func (m *middleware[M]) onErr(ctx context.Context, stage ErrorStage, err error) { + if err == nil { + return + } + if m.cfg != nil && m.cfg.OnError != nil { + m.cfg.OnError(ctx, stage, err) + } +} + +func (m *middleware[M]) lastUserMessage(agentIn *adk.TypedAgentInput[M]) (M, bool) { + if agentIn == nil || len(agentIn.Messages) == 0 { + return nil, false + } + if !m.topicSelectionEnabled() { + return nil, false + } + for i := len(agentIn.Messages) - 1; i >= 0; i-- { + msg := agentIn.Messages[i] + if isNilMessage(msg) || !isUserRole(msg) || isAutomemoryReminderMessage(msg) { + continue + } + return msg, true + } + return nil, false +} + +func (m *middleware[M]) topicSelectionTopK() int { + topK := m.cfg.Read.TopicSelection.TopK + if topK <= 0 { + return defaultTopicTopK + } + return topK +} + +func (m *middleware[M]) resolveSessionID(ctx context.Context, state *adk.TypedChatModelAgentState[M]) (string, error) { + if m.coordination != nil { + return strings.TrimSpace(m.coordination.SessionID), nil + } + return "", nil +} + +func (m *middleware[M]) sendTopicMemoryEvent(ctx context.Context, msgs []M, memMsg M) { + var beforeID string + if len(msgs) > 0 && !isNilMessage(msgs[len(msgs)-1]) { + beforeID = adk.GetMessageID(msgs[len(msgs)-1]) + } + if sendEventErr := adk.TypedSendEvent(ctx, &adk.TypedAgentEvent[M]{ + SessionEventVariant: &adk.SessionEventVariant[M]{ + Event: &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessageInserted, + MessageInserted: &adk.MessageInsertedEvent[M]{ + Message: memMsg, + BeforeMessageID: beforeID, + }, + }, + }, + }); sendEventErr != nil { + m.onErr(ctx, OnErrorStageSendSessionEvent, sendEventErr) + } +} diff --git a/adk/middlewares/backgroundtask/middleware.go b/adk/middlewares/backgroundtask/middleware.go new file mode 100644 index 000000000..bc4cf85c3 --- /dev/null +++ b/adk/middlewares/backgroundtask/middleware.go @@ -0,0 +1,320 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package backgroundtask provides the middleware that injects the background-task +// control tools (task_output, task_stop) into an agent. +// +// It is the single owner of these control tools: domain middlewares (subagent, +// filesystem) that launch background work register that work into a shared +// *backgroundtask.Manager, but they must NOT inject task_output/task_stop +// themselves. Wire this middleware exactly once per agent, bound to the same +// Manager the domain middlewares share, so the control tools are not duplicated. +package backgroundtask + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/cloudwego/eino/adk" + bgtask "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/schema" +) + +const ( + taskOutputToolName = "task_output" + taskStopToolName = "task_stop" +) + +// ToolConfig configures one of the injected control tools (task_output, task_stop). +type ToolConfig struct { + // Name overrides the tool name used in registration. + // Optional; the default name ("task_output" / "task_stop") is used when empty. + Name string + + // Desc overrides the tool description used in registration. + // Optional; the built-in description (with i18n) is used when nil. + Desc *string + + // Disable removes this tool from the injected set. + // Optional; false by default. Use it to expose only one of the control tools. + Disable bool +} + +// Config configures the background-task control middleware for the standard +// *schema.Message message type. It is the default specialization of TypedConfig. +type Config = TypedConfig[*schema.Message] + +// TypedConfig configures the background-task control middleware, parameterized by +// message type. +type TypedConfig[M adk.MessageType] struct { + // Manager is the shared background-task Manager whose tasks the injected + // task_output/task_stop tools inspect and cancel. Required. + // + // It is typically the same Manager the domain middlewares (subagent, filesystem) + // were given, so a single task-ID space spans agent and shell runs. + Manager *bgtask.Manager + + // TaskOutputToolConfig configures the task_output tool. Optional. + TaskOutputToolConfig *ToolConfig + // TaskStopToolConfig configures the task_stop tool. Optional. + TaskStopToolConfig *ToolConfig +} + +// New creates a middleware that injects the task_output and task_stop tools, bound +// to the Manager in config, for the standard *schema.Message message type. +func New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped[*schema.Message](ctx, config) +} + +// NewTyped creates a background-task control middleware parameterized by message type. +// See New for behavior details. +func NewTyped[M adk.MessageType](_ context.Context, config *TypedConfig[M]) (adk.TypedChatModelAgentMiddleware[M], error) { + if config == nil || config.Manager == nil { + return nil, fmt.Errorf("backgroundtask: Manager is required") + } + mgr := config.Manager + queried := newQueryTracker() + + outputEnabled := !disabled(config.TaskOutputToolConfig) + stopEnabled := !disabled(config.TaskStopToolConfig) + + var tools []tool.BaseTool + if outputEnabled { + outputTool, err := newTaskOutputTool(mgr, queried, config.TaskOutputToolConfig) + if err != nil { + return nil, fmt.Errorf("backgroundtask: failed to create task_output tool: %w", err) + } + tools = append(tools, outputTool) + } + if stopEnabled { + stopTool, err := newTaskStopTool(mgr, config.TaskStopToolConfig) + if err != nil { + return nil, fmt.Errorf("backgroundtask: failed to create task_stop tool: %w", err) + } + tools = append(tools, stopTool) + } + + instruction := buildInstruction(config.TaskOutputToolConfig, outputEnabled, config.TaskStopToolConfig, stopEnabled) + + return &typedMiddleware[M]{ + tools: tools, + instruction: instruction, + }, nil +} + +// disabled reports whether a tool config opts out of registering its tool. +func disabled(c *ToolConfig) bool { + return c != nil && c.Disable +} + +// buildInstruction assembles the background-task instruction so the per-tool +// sentences name the tools as actually registered and omit any disabled tool. +// It returns "" when no control tool is enabled, so a fully-disabled middleware +// injects nothing. +func buildInstruction(outputCfg *ToolConfig, outputEnabled bool, stopCfg *ToolConfig, stopEnabled bool) string { + if !outputEnabled && !stopEnabled { + return "" + } + + instruction := internal.SelectPrompt(internal.I18nPrompts{ + English: backgroundTaskPromptHeader, + Chinese: backgroundTaskPromptHeaderChinese, + }) + if outputEnabled { + line := internal.SelectPrompt(internal.I18nPrompts{ + English: backgroundTaskOutputLine, + Chinese: backgroundTaskOutputLineChinese, + }) + instruction += fmt.Sprintf(line, selectToolName(outputCfg, taskOutputToolName)) + } + if stopEnabled { + line := internal.SelectPrompt(internal.I18nPrompts{ + English: backgroundTaskStopLine, + Chinese: backgroundTaskStopLineChinese, + }) + instruction += fmt.Sprintf(line, selectToolName(stopCfg, taskStopToolName)) + } + instruction += internal.SelectPrompt(internal.I18nPrompts{ + English: backgroundTaskPromptFooter, + Chinese: backgroundTaskPromptFooterChinese, + }) + return instruction +} + +// selectToolName returns the configured name override, or the default when unset. +func selectToolName(c *ToolConfig, defaultName string) string { + if c != nil && c.Name != "" { + return c.Name + } + return defaultName +} + +// selectToolDesc returns the configured description override, or the built-in +// i18n description when unset. +func selectToolDesc(c *ToolConfig, english, chinese string) string { + if c != nil && c.Desc != nil { + return *c.Desc + } + return internal.SelectPrompt(internal.I18nPrompts{English: english, Chinese: chinese}) +} + +type typedMiddleware[M adk.MessageType] struct { + adk.TypedBaseChatModelAgentMiddleware[M] + tools []tool.BaseTool + instruction string +} + +// BeforeAgent injects the control tools and instruction into the agent context. +func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { + if runCtx == nil { + return ctx, runCtx, nil + } + + nRunCtx := *runCtx + if m.instruction != "" { + nRunCtx.Instruction += "\n" + m.instruction + } + nRunCtx.Tools = append(nRunCtx.Tools, m.tools...) + return ctx, &nRunCtx, nil +} + +type taskOutputInput struct { + TaskID string `json:"task_id" jsonschema:"required" jsonschema_description:"The task ID to get output from"` + // Block defaults to true (wait for the task to finish). A *bool distinguishes + // "omitted" (wait) from an explicit false (return the current status now). + Block *bool `json:"block,omitempty" jsonschema_description:"Whether to wait for the task to complete. Defaults to true; set false to return the current status immediately."` + Timeout int `json:"timeout,omitempty" jsonschema_description:"Maximum time to wait in milliseconds when blocking. Defaults to 30000; capped at 600000."` +} + +// queryTracker is owned by the task_output middleware, not by Manager. It keeps +// consumption bookkeeping out of the lifecycle registry. +type queryTracker struct { + mu sync.Mutex + queried map[string]struct{} +} + +func newQueryTracker() *queryTracker { + return &queryTracker{queried: make(map[string]struct{})} +} + +func (q *queryTracker) mark(id string) { + q.mu.Lock() + defer q.mu.Unlock() + q.queried[id] = struct{}{} +} + +const ( + defaultTaskOutputTimeoutMs = 30000 + maxTaskOutputTimeoutMs = 600000 +) + +func newTaskOutputTool(mgr *bgtask.Manager, queried *queryTracker, cfg *ToolConfig) (tool.InvokableTool, error) { + name := selectToolName(cfg, taskOutputToolName) + desc := selectToolDesc(cfg, taskOutputToolDescription, taskOutputToolDescriptionChinese) + return utils.InferTool(name, desc, func(ctx context.Context, input taskOutputInput) (string, error) { + task, ok := resolveTask(ctx, mgr, input) + if !ok { + return fmt.Sprintf("Task %q not found", input.TaskID), nil + } + + // Only mark the result as consumed once the task has actually finished. + // A still-running task has no final result yet, so polling its status + // must not mark a never-read result as consumed. + if task.Status != bgtask.StatusRunning { + queried.mark(input.TaskID) + } + + return formatTask(task), nil + }) +} + +// resolveTask fetches the task, optionally blocking until it finishes. Blocking is +// the default; it is bounded by input.Timeout (clamped to [0, max], default 30s). +// The returned bool reports whether the task exists (not whether it finished). +func resolveTask(ctx context.Context, mgr *bgtask.Manager, input taskOutputInput) (*bgtask.Task, bool) { + if input.Block != nil && !*input.Block { + return mgr.Get(input.TaskID) + } + + timeoutMs := input.Timeout + if timeoutMs <= 0 { + timeoutMs = defaultTaskOutputTimeoutMs + } + if timeoutMs > maxTaskOutputTimeoutMs { + timeoutMs = maxTaskOutputTimeoutMs + } + + waitCtx, cancel := context.WithTimeout(ctx, time.Duration(timeoutMs)*time.Millisecond) + defer cancel() + // Wait's bool reports whether the task reached a terminal state; for the tool we + // only care whether the task exists, so translate via the returned snapshot. + task, _ := mgr.Wait(waitCtx, input.TaskID) + return task, task != nil +} + +type taskStopInput struct { + TaskID string `json:"task_id" jsonschema:"required" jsonschema_description:"The ID of the background task to stop"` +} + +func newTaskStopTool(mgr *bgtask.Manager, cfg *ToolConfig) (tool.InvokableTool, error) { + name := selectToolName(cfg, taskStopToolName) + desc := selectToolDesc(cfg, taskStopToolDescription, taskStopToolDescriptionChinese) + return utils.InferTool(name, desc, func(ctx context.Context, input taskStopInput) (string, error) { + if err := mgr.Cancel(input.TaskID); err != nil { + return fmt.Sprintf("Failed to stop task %q: %s", input.TaskID, err.Error()), nil + } + return fmt.Sprintf("Successfully stopped task: %s", input.TaskID), nil + }) +} + +func formatTask(task *bgtask.Task) string { + result := fmt.Sprintf("Task ID: %s\nDescription: %s\nStatus: %s", + task.ID, task.Description, task.Status) + + // When the task has a reliable output file, the file is authoritative — point at + // it and do not inline Result. The file carries the same (or interim) output and + // may be large, so Read'ing it selectively avoids inlining the whole blob. When a + // write to the file failed (OutputFileErr set), neither side is the complete + // output: the file has a gap, and Result is only what the worker returned (which + // may be empty while the task runs, or a partial projection of the file). Report + // the failure honestly and surface Result as best-effort current data rather than + // presenting either as authoritative. Without an output file, Result is the only + // copy, so inline it. + if task.OutputFile != "" && task.OutputFileErr == "" { + result += fmt.Sprintf("\nOutput file: %s (use Read on this path for the output)", task.OutputFile) + } else { + if task.Result != "" { + result += fmt.Sprintf("\nResult: %s", task.Result) + } + if task.OutputFile != "" { + result += fmt.Sprintf("\nOutput file: %s (incomplete — a write failed: %s; full output is unavailable. The Result above, if any, is the best-effort output captured so far and may be empty or partial)", + task.OutputFile, task.OutputFileErr) + } + } + if task.Error != "" { + result += fmt.Sprintf("\nError: %s", task.Error) + } + if task.DoneAt != nil { + result += fmt.Sprintf("\nCompleted at: %s", task.DoneAt.Format("2006-01-02 15:04:05")) + } + + return result +} diff --git a/adk/middlewares/backgroundtask/middleware_test.go b/adk/middlewares/backgroundtask/middleware_test.go new file mode 100644 index 000000000..d3fa32be4 --- /dev/null +++ b/adk/middlewares/backgroundtask/middleware_test.go @@ -0,0 +1,323 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package backgroundtask + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + bgtask "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func closeWithTimeout(m *bgtask.Manager) { + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + _ = m.Close(ctx) +} + +func runWork(m *bgtask.Manager, description string, background bool, work bgtask.WorkFunc) (*bgtask.Task, error) { + return m.Run(context.Background(), &bgtask.RunInput{ + Description: description, + RunInBackground: background, + }, work) +} + +func completedWork(result string) bgtask.WorkFunc { + return func(ctx context.Context, _ bgtask.TaskInfo) (string, error) { + return result, nil + } +} + +func blockingWork() bgtask.WorkFunc { + return func(ctx context.Context, _ bgtask.TaskInfo) (string, error) { + <-ctx.Done() + return "", ctx.Err() + } +} + +// findTool returns the named tool from a tool list. +func findTool(t *testing.T, tools []tool.BaseTool, name string) tool.InvokableTool { + t.Helper() + for _, bt := range tools { + info, err := bt.Info(context.Background()) + require.NoError(t, err) + if info.Name == name { + it, ok := bt.(tool.InvokableTool) + require.True(t, ok) + return it + } + } + t.Fatalf("tool %q not found", name) + return nil +} + +func injectedTools(t *testing.T, m *bgtask.Manager) []tool.BaseTool { + t.Helper() + mw, err := New(context.Background(), &Config{Manager: m}) + require.NoError(t, err) + _, runCtx, err := mw.BeforeAgent(context.Background(), &adk.ChatModelAgentContext[*schema.Message]{}) + require.NoError(t, err) + return runCtx.Tools +} + +func TestNew_NilManager(t *testing.T) { + _, err := New(context.Background(), nil) + assert.Error(t, err) +} + +func TestMiddleware_InjectsControlTools(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + tools := injectedTools(t, mgr) + require.Len(t, tools, 2) + + // Both control tools present. + findTool(t, tools, taskOutputToolName) + findTool(t, tools, taskStopToolName) +} + +func TestMiddleware_ToolConfig_NameOverrideAndDisable(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + customDesc := "custom output desc" + mw, err := New(context.Background(), &Config{ + Manager: mgr, + TaskOutputToolConfig: &ToolConfig{Name: "get_output", Desc: &customDesc}, + TaskStopToolConfig: &ToolConfig{Disable: true}, + }) + require.NoError(t, err) + _, runCtx, err := mw.BeforeAgent(context.Background(), &adk.ChatModelAgentContext[*schema.Message]{}) + require.NoError(t, err) + + // task_stop disabled → only the renamed task_output remains. + require.Len(t, runCtx.Tools, 1) + info, err := runCtx.Tools[0].Info(context.Background()) + require.NoError(t, err) + assert.Equal(t, "get_output", info.Name) + assert.Equal(t, customDesc, info.Desc) +} + +func TestMiddleware_ToolConfig_DisableBoth(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + mw, err := New(context.Background(), &Config{ + Manager: mgr, + TaskOutputToolConfig: &ToolConfig{Disable: true}, + TaskStopToolConfig: &ToolConfig{Disable: true}, + }) + require.NoError(t, err) + _, runCtx, err := mw.BeforeAgent(context.Background(), &adk.ChatModelAgentContext[*schema.Message]{}) + require.NoError(t, err) + assert.Empty(t, runCtx.Tools) +} + +func TestMiddleware_InjectsInstruction(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + mw, err := New(context.Background(), &Config{Manager: mgr}) + require.NoError(t, err) + _, runCtx, err := mw.BeforeAgent(context.Background(), &adk.ChatModelAgentContext[*schema.Message]{Instruction: "base"}) + require.NoError(t, err) + assert.Contains(t, runCtx.Instruction, "base") + assert.Contains(t, runCtx.Instruction, "task_output") + assert.Contains(t, runCtx.Instruction, "task_stop") +} + +// TestMiddleware_InstructionUsesRenamedTool verifies the instruction names the +// tool as registered: a renamed task_output is referenced by its new name, and +// the default name no longer appears. +func TestMiddleware_InstructionUsesRenamedTool(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + mw, err := New(context.Background(), &Config{ + Manager: mgr, + TaskOutputToolConfig: &ToolConfig{Name: "get_task_result"}, + }) + require.NoError(t, err) + _, runCtx, err := mw.BeforeAgent(context.Background(), &adk.ChatModelAgentContext[*schema.Message]{}) + require.NoError(t, err) + assert.Contains(t, runCtx.Instruction, "get_task_result") + assert.NotContains(t, runCtx.Instruction, "task_output") + assert.Contains(t, runCtx.Instruction, "task_stop") +} + +// TestMiddleware_InstructionOmitsDisabledTool verifies a disabled tool's sentence +// is dropped so the model is never told to call a tool that was not registered. +func TestMiddleware_InstructionOmitsDisabledTool(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + mw, err := New(context.Background(), &Config{ + Manager: mgr, + TaskStopToolConfig: &ToolConfig{Disable: true}, + }) + require.NoError(t, err) + _, runCtx, err := mw.BeforeAgent(context.Background(), &adk.ChatModelAgentContext[*schema.Message]{}) + require.NoError(t, err) + assert.Contains(t, runCtx.Instruction, "task_output") + assert.NotContains(t, runCtx.Instruction, "task_stop") +} + +// TestMiddleware_InstructionEmptyWhenAllDisabled verifies a fully-disabled +// middleware injects neither tools nor a background-task instruction. +func TestMiddleware_InstructionEmptyWhenAllDisabled(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + mw, err := New(context.Background(), &Config{ + Manager: mgr, + TaskOutputToolConfig: &ToolConfig{Disable: true}, + TaskStopToolConfig: &ToolConfig{Disable: true}, + }) + require.NoError(t, err) + _, runCtx, err := mw.BeforeAgent(context.Background(), &adk.ChatModelAgentContext[*schema.Message]{Instruction: "base"}) + require.NoError(t, err) + assert.Equal(t, "base", runCtx.Instruction) + assert.Empty(t, runCtx.Tools) +} + +func TestTaskOutputTool(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + result, err := runWork(mgr, "test task", false, completedWork("task result")) + require.NoError(t, err) + require.Equal(t, bgtask.StatusCompleted, result.Status) + + tl := findTool(t, injectedTools(t, mgr), taskOutputToolName) + output, err := tl.InvokableRun(context.Background(), fmt.Sprintf(`{"task_id":"%s"}`, result.ID)) + require.NoError(t, err) + assert.Contains(t, output, "test task") + assert.Contains(t, output, "task result") + assert.Contains(t, output, "completed") +} + +func TestTaskOutputTool_NotFound(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + tl := findTool(t, injectedTools(t, mgr), taskOutputToolName) + result, err := tl.InvokableRun(context.Background(), `{"task_id":"nonexistent"}`) + require.NoError(t, err) + assert.Contains(t, result, "not found") +} + +func TestTaskOutputTool_NonBlockingRunningThenTerminal(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + runResult, err := runWork(mgr, "running task", true, blockingWork()) + require.NoError(t, err) + + tl := findTool(t, injectedTools(t, mgr), taskOutputToolName) + out, err := tl.InvokableRun(context.Background(), fmt.Sprintf(`{"task_id":"%s","block":false}`, runResult.ID)) + require.NoError(t, err) + assert.Contains(t, out, "running") + + require.NoError(t, mgr.Cancel(runResult.ID)) + waitCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + task, done := mgr.Wait(waitCtx, runResult.ID) + require.True(t, done) + require.NotNil(t, task) + + _, err = tl.InvokableRun(context.Background(), fmt.Sprintf(`{"task_id":"%s","block":false}`, runResult.ID)) + require.NoError(t, err) +} + +func TestTaskStopTool(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + runResult, err := runWork(mgr, "running task", true, blockingWork()) + require.NoError(t, err) + + tl := findTool(t, injectedTools(t, mgr), taskStopToolName) + result, err := tl.InvokableRun(context.Background(), fmt.Sprintf(`{"task_id":"%s"}`, runResult.ID)) + require.NoError(t, err) + assert.Contains(t, result, "Successfully stopped") + + task, ok := mgr.Get(runResult.ID) + require.True(t, ok) + assert.Equal(t, bgtask.StatusCanceled, task.Status) +} + +func TestTaskStopTool_AlreadyDone(t *testing.T) { + mgr := bgtask.New(context.Background(), &bgtask.Config{}) + defer closeWithTimeout(mgr) + + runResult, err := runWork(mgr, "done task", false, completedWork("done")) + require.NoError(t, err) + require.Equal(t, bgtask.StatusCompleted, runResult.Status) + + tl := findTool(t, injectedTools(t, mgr), taskStopToolName) + result, err := tl.InvokableRun(context.Background(), fmt.Sprintf(`{"task_id":"%s"}`, runResult.ID)) + require.NoError(t, err) + assert.Contains(t, result, "Failed to stop") +} + +// A reliable output file is authoritative: formatTask points at it and does not +// inline Result. +func TestFormatTask_ReliableOutputFile(t *testing.T) { + out := formatTask(&bgtask.Task{ + ID: "bash_1", + Status: bgtask.StatusCompleted, + Result: "the full result", + OutputFile: "/tasks/bash_1.output", + }) + assert.Contains(t, out, "/tasks/bash_1.output") + assert.NotContains(t, out, "the full result", "a reliable file replaces inlining Result") +} + +// When the output file is marked unreliable, formatTask falls back to the complete +// in-memory Result and flags the file as incomplete rather than pointing at it as +// the sole authority. +func TestFormatTask_UnreliableOutputFile_FallsBackToResult(t *testing.T) { + out := formatTask(&bgtask.Task{ + ID: "bash_1", + Status: bgtask.StatusCompleted, + Result: "the full result", + OutputFile: "/tasks/bash_1.output", + OutputFileErr: "append failed", + }) + assert.Contains(t, out, "the full result", "Result must be surfaced when the file is unreliable") + assert.Contains(t, out, "incomplete", "the file must be flagged as partial") + assert.Contains(t, out, "/tasks/bash_1.output") +} + +// With no output file, Result is the only copy and is inlined. +func TestFormatTask_NoOutputFile(t *testing.T) { + out := formatTask(&bgtask.Task{ + ID: "bash_1", + Status: bgtask.StatusCompleted, + Result: "the full result", + }) + assert.Contains(t, out, "the full result") +} diff --git a/adk/middlewares/backgroundtask/prompt.go b/adk/middlewares/backgroundtask/prompt.go new file mode 100644 index 000000000..a8d77a29d --- /dev/null +++ b/adk/middlewares/backgroundtask/prompt.go @@ -0,0 +1,77 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package backgroundtask + +const ( + taskOutputToolDescription = `Retrieve the output and status of a running or completed background task. + +- Takes a task_id parameter identifying the task +- Returns the task's output along with its status and any error +- Use this tool to check on a background task or retrieve its result by task_id +` + + taskOutputToolDescriptionChinese = `获取正在运行或已完成的后台任务的输出与状态。 + +- 接受 task_id 参数来标识任务 +- 返回该任务的输出，以及其状态和任何错误信息 +- 当你需要查询后台任务或通过 task_id 获取其结果时使用此工具 +` + + taskStopToolDescription = `Stop a running background task by its ID. + +- Takes a task_id parameter identifying the task to stop +- Returns a success or failure status +- Use this tool when you need to cancel a long-running background task +` + + taskStopToolDescriptionChinese = `通过 ID 停止正在运行的后台任务。 + +- 接受 task_id 参数来标识要停止的任务 +- 返回成功或失败状态 +- 当你需要取消一个长时间运行的后台任务时使用此工具 +` + + // The instruction is assembled from these pieces so the per-tool sentences name + // the tools as actually registered: a tool renamed via ToolConfig.Name is + // referenced by that name, and a disabled tool's sentence is omitted entirely. + // Keeping the model's instructions in sync with the live tool set avoids telling + // it to call a tool that was renamed or no longer exists. + backgroundTaskPromptHeader = ` +## Background Task Management +- Some tools can launch work in the background. Background tasks keep running after the + tool call returns; you will be notified when they complete.` + + // %s is the registered task_output tool name. + backgroundTaskOutputLine = "\n- Use the %s tool to check a background task's status or retrieve its result by task_id." + + // %s is the registered task_stop tool name. + backgroundTaskStopLine = "\n- Use the %s tool to cancel a running background task by task_id." + + backgroundTaskPromptFooter = "\n- These tasks are running executions, not planning to-dos.\n" + + backgroundTaskPromptHeaderChinese = ` +## 后台任务管理 +- 部分工具可以在后台启动任务。后台任务在工具调用返回后会继续运行；任务完成时你将收到通知。` + + // %s is the registered task_output tool name. + backgroundTaskOutputLineChinese = "\n- 使用 %s 工具通过 task_id 查询后台任务的状态或获取其结果。" + + // %s is the registered task_stop tool name. + backgroundTaskStopLineChinese = "\n- 使用 %s 工具通过 task_id 取消正在运行的后台任务。" + + backgroundTaskPromptFooterChinese = "\n- 这些任务是正在运行的执行实例，而非用于规划的待办事项。\n" +) diff --git a/adk/middlewares/dynamictool/toolsearch/toolsearch.go b/adk/middlewares/dynamictool/toolsearch/toolsearch.go index 9215b1964..b2948382c 100644 --- a/adk/middlewares/dynamictool/toolsearch/toolsearch.go +++ b/adk/middlewares/dynamictool/toolsearch/toolsearch.go @@ -134,7 +134,7 @@ type typedMiddleware[M adk.MessageType] struct { sr string } -func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { +func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { if runCtx == nil { return ctx, runCtx, nil } @@ -155,38 +155,46 @@ const toolSearchReminderExtraKey = "__toolsearch_reminder__" func (m *typedMiddleware[M]) isInitialized(ctx context.Context) bool { val, ok, err := adk.GetRunLocalValue(ctx, toolSearchInitializedKey) - if err != nil || !ok { - return false + if err == nil && ok { + if b, _ := val.(bool); b { + return true + } } - b, _ := val.(bool) - return b + return false } func (m *typedMiddleware[M]) markInitialized(ctx context.Context) { _ = adk.SetRunLocalValue(ctx, toolSearchInitializedKey, true) } -func (m *typedMiddleware[M]) ensureReminder(msgs []M) []M { +func (m *typedMiddleware[M]) ensureReminder(msgs []M) (result []M, insertedMsg M, anchorMsg M, didInsert bool) { for _, msg := range msgs { if hasToolSearchReminderExtra(msg) { - return msgs + return msgs, insertedMsg, anchorMsg, false } } - reminder := makeReminderMsg[M](m.sr) - result := make([]M, 0, len(msgs)+1) + insertedMsg = makeReminderMsg[M](m.sr) + adk.EnsureMessageID(insertedMsg) + result = make([]M, 0, len(msgs)+1) inserted := false for _, msg := range msgs { if !inserted && !isSystemRoleTS(msg) { inserted = true - result = append(result, reminder) + result = append(result, insertedMsg) + anchorMsg = msg } result = append(result, msg) } if !inserted { - result = append(result, reminder) + result = append(result, insertedMsg) } - return result + return result, insertedMsg, anchorMsg, true +} + +func isNilTSMessage[M adk.MessageType](msg M) bool { + var zero M + return any(msg) == any(zero) } func isSystemRoleTS[M adk.MessageType](msg M) bool { @@ -275,7 +283,26 @@ func toolNameSet(tools []*schema.ToolInfo) map[string]bool { } func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) { - state.Messages = m.ensureReminder(state.Messages) + newMsgs, insertedMsg, anchorMsg, didInsert := m.ensureReminder(state.Messages) + state.Messages = newMsgs + + if didInsert { + var beforeID string + if !isNilTSMessage(anchorMsg) { + beforeID = adk.GetMessageID(anchorMsg) + } + _ = adk.TypedSendEvent(ctx, &adk.TypedAgentEvent[M]{ + SessionEventVariant: &adk.SessionEventVariant[M]{ + Event: &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessageInserted, + MessageInserted: &adk.MessageInsertedEvent[M]{ + Message: insertedMsg, + BeforeMessageID: beforeID, + }, + }, + }, + }) + } if !m.isInitialized(ctx) { m.markInitialized(ctx) diff --git a/adk/middlewares/dynamictool/toolsearch/toolsearch_generic_test.go b/adk/middlewares/dynamictool/toolsearch/toolsearch_generic_test.go index a659f07df..76ecd89b9 100644 --- a/adk/middlewares/dynamictool/toolsearch/toolsearch_generic_test.go +++ b/adk/middlewares/dynamictool/toolsearch/toolsearch_generic_test.go @@ -213,7 +213,7 @@ func testEnsureReminderGeneric[M adk.MessageType](t *testing.T) { makeSystemMsg[M]("sys"), makeUserMsg[M]("hi"), } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 3) assert.Equal(t, "system", getMsgRole(got[0])) // Reminder inserted after system @@ -228,7 +228,7 @@ func testEnsureReminderGeneric[M adk.MessageType](t *testing.T) { makeSystemMsg[M]("sys1"), makeSystemMsg[M]("sys2"), } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 3) assert.Equal(t, "system", getMsgRole(got[0])) assert.Equal(t, "system", getMsgRole(got[1])) @@ -239,7 +239,7 @@ func testEnsureReminderGeneric[M adk.MessageType](t *testing.T) { }) t.Run("empty input", func(t *testing.T) { - got := m.ensureReminder(nil) + got, _, _, _ := m.ensureReminder(nil) require.Len(t, got, 1) extra := getMsgExtra(got[0]) require.NotNil(t, extra) @@ -250,7 +250,7 @@ func testEnsureReminderGeneric[M adk.MessageType](t *testing.T) { input := []M{ makeUserMsg[M]("hi"), } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 2) // Reminder inserted at position 0 extra := getMsgExtra(got[0]) @@ -266,7 +266,7 @@ func testEnsureReminderGeneric[M adk.MessageType](t *testing.T) { reminder, makeUserMsg[M]("hi"), } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 2) assert.Equal(t, "hi", getMsgContent(got[1])) }) diff --git a/adk/middlewares/dynamictool/toolsearch/toolsearch_test.go b/adk/middlewares/dynamictool/toolsearch/toolsearch_test.go index 4bd1410ec..789f902c9 100644 --- a/adk/middlewares/dynamictool/toolsearch/toolsearch_test.go +++ b/adk/middlewares/dynamictool/toolsearch/toolsearch_test.go @@ -446,7 +446,7 @@ func TestEnsureReminder(t *testing.T) { {Role: schema.System, Content: "sys"}, {Role: schema.User, Content: "hi"}, } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 3) assert.Equal(t, schema.System, got[0].Role) assert.Equal(t, schema.User, got[1].Role) @@ -461,7 +461,7 @@ func TestEnsureReminder(t *testing.T) { {Role: schema.System, Content: "sys1"}, {Role: schema.System, Content: "sys2"}, } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 3) assert.Equal(t, schema.System, got[0].Role) assert.Equal(t, schema.System, got[1].Role) @@ -469,7 +469,7 @@ func TestEnsureReminder(t *testing.T) { }) t.Run("empty input", func(t *testing.T) { - got := m.ensureReminder(nil) + got, _, _, _ := m.ensureReminder(nil) require.Len(t, got, 1) assert.Equal(t, "", got[0].Content) }) @@ -479,7 +479,7 @@ func TestEnsureReminder(t *testing.T) { {Role: schema.User, Content: "hi"}, {Role: schema.Assistant, Content: "hello"}, } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 3) assert.Equal(t, "", got[0].Content) assert.Equal(t, "hi", got[1].Content) @@ -491,7 +491,7 @@ func TestEnsureReminder(t *testing.T) { {Role: schema.User, Content: "", Extra: map[string]any{toolSearchReminderExtraKey: true}}, {Role: schema.User, Content: "hi"}, } - got := m.ensureReminder(input) + got, _, _, _ := m.ensureReminder(input) require.Len(t, got, 2) assert.Equal(t, "", got[0].Content) assert.Equal(t, "hi", got[1].Content) diff --git a/adk/middlewares/filesystem/bash_run.go b/adk/middlewares/filesystem/bash_run.go new file mode 100644 index 000000000..64a9cc8a8 --- /dev/null +++ b/adk/middlewares/filesystem/bash_run.go @@ -0,0 +1,310 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "context" + "fmt" + "io" + "path/filepath" + + "github.com/google/uuid" + + "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// ExecuteTaskType is the backgroundtask Task.Type tag for shell-command tasks +// launched by the execute tool. A shared Manager's ShouldAutoBackground hook can +// match on it to apply shell-specific policy, recovering the command via +// CommandFromTask. +// +// When the filesystem middleware is configured with both a Backend and an +// OutputDir, the managed execute tool writes each task's output to a file under +// that directory and records the path on Task.OutputFile (streaming runs tee +// chunks as interim output; buffered runs write the result on completion). The +// Manager itself never writes — the execute tool owns it. +const ExecuteTaskType = "bash" + +// MetadataKeyCommand is the RunInput.Metadata / Task.Metadata key under which the +// execute tool records the shell command for a task. A ShouldAutoBackground hook +// reads it (via CommandFromTask) to apply command-specific policy without parsing +// the human-readable Description. The value is a string. +const MetadataKeyCommand = "command" + +// CommandFromTask returns the shell command recorded in a shell task's metadata +// under MetadataKeyCommand. The execute tool always records it for shell tasks, so +// a hook receiving a task of ExecuteTaskType can rely on a non-empty result; it +// returns "" only when given a nil or non-shell task. +func CommandFromTask(t *backgroundtask.Task) string { + if t == nil { + return "" + } + cmd, _ := t.Metadata[MetadataKeyCommand].(string) + return cmd +} + +// outputSink bundles the output-file configuration for a managed execute tool: an +// Appender to write through and the directory to reserve paths under. Both must be +// set to enable output files; a zero outputSink disables them. +type outputSink struct { + appender filesystem.Appender + outputDir string +} + +// bashOutputWriter appends a managed execute task's output to a file via a +// filesystem.Appender. It is built per invocation: when both an Appender and an +// outputDir are configured it reserves outputDir/.output and appends there; +// otherwise it is disabled and every method is a no-op, so the task has no output +// file. There is no rewrite fallback — output files require an Appender. +// +// The execute tool — not the Manager — owns writing, so streaming runs tee interim +// output as chunks arrive. It is single-consumer: append is called serially on +// the StreamReaderWithConvert Recv stack, so no synchronization is needed. +// +// On the first append failure the file is left with a gap, so the writer records +// the failure via mgr.MarkOutputFileUnreliable (keyed by taskID, which the work +// func receives from the Manager and sets on the writer before its first append) +// and stops attempting further writes. +type bashOutputWriter struct { + mgr *backgroundtask.Manager + appender filesystem.Appender // nil => disabled + path string + taskID string // set by the work func once the Manager assigns it + failed bool // set after the first append error: the file is now partial +} + +// reserveBashOutput builds a writer that appends under the sink, or a disabled +// writer when output files are not configured (no appender / no dir). It creates +// the file empty up front so the advertised path exists before any output; if even +// that reservation write fails, it returns a disabled writer so the task advertises +// no output file and consumers fall back to the in-memory Result. The file is +// named after the launching tool-call id (so it matches Task.ToolUseID), falling +// back to a uuid when no tool-call id is in context. +func reserveBashOutput(ctx context.Context, mgr *backgroundtask.Manager, sink outputSink) *bashOutputWriter { + if sink.appender == nil || sink.outputDir == "" { + return &bashOutputWriter{} + } + path := filepath.Join(sink.outputDir, outputFileName(ctx)+".output") + if err := sink.appender.Append(ctx, &filesystem.AppendRequest{FilePath: path, Content: ""}); err != nil { + return &bashOutputWriter{} + } + return &bashOutputWriter{ + mgr: mgr, + appender: sink.appender, + path: path, + } +} + +func (w *bashOutputWriter) append(ctx context.Context, content string) { + if w.appender == nil || w.failed { + return + } + if err := w.appender.Append(ctx, &filesystem.AppendRequest{FilePath: w.path, Content: content}); err != nil { + // The file now has a gap: stop writing and mark it unreliable (by task id) so + // task_output reports the file's failed state instead of trusting the partial file. + w.failed = true + w.mgr.MarkOutputFileUnreliable(w.taskID, err.Error()) + } +} + +// outputFileName returns the base name (without extension) for a task's output +// file: the launching tool-call id when present (so the file matches +// Task.ToolUseID), or a uuid fallback when no tool-call id is in context — the +// fallback keeps names unique so concurrent untagged tasks don't collide. +func outputFileName(ctx context.Context) string { + if id := compose.GetToolCallID(ctx); id != "" { + return id + } + return uuid.NewString() +} + +// bashWork adapts a blocking shell execution into a backgroundtask.WorkFunc. +// The request carries only the command; the Manager is the sole owner of +// foreground/background/auto-background switching, so no background hint is +// pushed down to the backend. On success it appends the result to the output file +// (when one is configured) before returning, so the file matches Task.Result. +func bashWork(sb filesystem.Shell, req *filesystem.ExecuteRequest, w *bashOutputWriter) backgroundtask.WorkFunc { + return func(ctx context.Context, task backgroundtask.TaskInfo) (string, error) { + w.taskID = task.ID + result, err := sb.Execute(ctx, req) + if err != nil { + return "", err + } + out := convExecuteResponse(result) + w.append(ctx, out) + return out, nil + } +} + +// bashStreamWork adapts a streaming shell execution into a backgroundtask.StreamWorkFunc. +// It returns a stream of formatted output chunks; the Manager forwards them to the +// caller in real time (for the foreground phase) and accumulates them into the +// task's final result. The terminal note (exit code / no-output) is emitted as a +// final chunk so it is part of both the live stream and the persisted result. +// +// Each emitted chunk (and the terminal note) is also teed to the output file via w, +// so the file carries interim output while the task runs. Teeing happens inside the +// convert/OnEOF callbacks, which run on the Recv stack for both the foreground loop +// and the background drain — so the Manager never has to write. +func bashStreamWork(sb filesystem.StreamingShell, req *filesystem.ExecuteRequest, w *bashOutputWriter) backgroundtask.StreamWorkFunc { + return func(ctx context.Context, task backgroundtask.TaskInfo) (*schema.StreamReader[string], error) { + w.taskID = task.ID + stream, err := sb.ExecuteStreaming(ctx, req) + if err != nil { + return nil, err + } + + // exitCode/hasContent accumulate across chunks: convert writes them per + // chunk, the OnEOF hook reads them to build the terminal note. The convert + // model has no per-stream state of its own, so they live in this closure. + // Safe without synchronization because StreamReaderWithConvert is pull-driven + // and single-consumer — convert and OnEOF run serially on the same Recv stack. + var exitCode *int + var hasContent bool + return schema.StreamReaderWithConvert(stream, + func(chunk *filesystem.ExecuteResponse) (string, error) { + if chunk == nil { + return "", schema.ErrNoValue + } + if chunk.ExitCode != nil { + exitCode = chunk.ExitCode + } + text := formatExecChunk(chunk.Output, chunk.Truncated) + if text == "" { + return "", schema.ErrNoValue + } + hasContent = true + w.append(ctx, text) + return text, nil + }, + schema.WithOnEOF(func() (any, error) { + if note := execTerminalNote(exitCode, hasContent); note != "" { + w.append(ctx, note) + return note, nil + } + return nil, io.EOF + }), + ), nil + } +} + +// newManagedExecuteTool builds an execute tool whose runs are tracked by a shared +// background-task Manager. The model controls background execution via the +// run_in_background field; auto-background switching is handled transparently by +// the Manager. On a background launch the tool returns the task ID so the agent +// can later query it via task_output. +// +// With a StreamingShell backend the tool is itself a StreamableTool: the +// foreground phase streams output to the caller in real time, and a run that moves +// to the background caps the stream with a notice (the rest is drained into the +// task result). With a plain Shell backend the tool is buffered. +// +// Exactly one of sb / streaming must be non-nil. appender and outputDir, when both +// set, enable per-task output files (the tool appends output to +// outputDir/.output); otherwise runs have no output file. +// Exactly one of sb / streaming must be non-nil. sink, when fully configured +// (appender + dir), enables per-task output files (the tool appends output to +// outputDir/.output); otherwise runs have no output file. +func newManagedExecuteTool( + mgr *backgroundtask.Manager, + sb filesystem.Shell, + streaming filesystem.StreamingShell, + sink outputSink, + name string, + desc string, +) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameExecute) + d, err := selectToolDesc(desc, ManagedExecuteToolDesc, ManagedExecuteToolDescChinese) + if err != nil { + return nil, err + } + + if streaming != nil { + return newManagedStreamingExecuteTool(mgr, streaming, sink, toolName, d) + } + return newManagedBufferedExecuteTool(mgr, sb, sink, toolName, d) +} + +// managedRunInput builds the RunInput shared by the buffered and streaming managed +// execute tools. w supplies the reserved output-file path (empty when output files +// are not configured), which the work funcs write to. +func managedRunInput(ctx context.Context, input executeManagedArgs, w *bashOutputWriter) *backgroundtask.RunInput { + runInput := &backgroundtask.RunInput{ + Description: input.Command, + Type: ExecuteTaskType, + ToolUseID: compose.GetToolCallID(ctx), + RunInBackground: input.RunInBackground, + Metadata: map[string]any{MetadataKeyCommand: input.Command}, + OutputFile: w.path, + } + // A positive timeout overrides the Manager's default foreground budget for + // this command. When the deadline expires, the Manager's policy decides + // whether to move the task to the background or stop it. + if input.TimeoutMS > 0 { + runInput.ForegroundTimeoutMs = &input.TimeoutMS + } + return runInput +} + +func newManagedBufferedExecuteTool(mgr *backgroundtask.Manager, sb filesystem.Shell, sink outputSink, toolName, desc string) (tool.BaseTool, error) { + return utils.InferTool(toolName, desc, func(ctx context.Context, input executeManagedArgs) (string, error) { + req := &filesystem.ExecuteRequest{Command: input.Command} + w := reserveBashOutput(ctx, mgr, sink) + result, err := mgr.Run(ctx, managedRunInput(ctx, input, w), bashWork(sb, req, w)) + if err != nil { + return "", err + } + + switch result.Status { + case backgroundtask.StatusCompleted: + return result.Result, nil + case backgroundtask.StatusRunning: + msg := fmt.Sprintf("Command running in background with ID: %s.", result.ID) + if result.OutputFile != "" { + msg += fmt.Sprintf(" Output is being written to: %s.", result.OutputFile) + } + msg += " You will be notified when it completes." + if result.OutputFile != "" { + msg += " To check interim output, use Read on that file path." + } + return msg, nil + case backgroundtask.StatusFailed: + return "", fmt.Errorf("execute task %q failed: %s", result.ID, result.Error) + case backgroundtask.StatusCanceled: + return "", fmt.Errorf("execute task %q was canceled", result.ID) + default: + return result.Result, nil + } + }) +} + +func newManagedStreamingExecuteTool(mgr *backgroundtask.Manager, streaming filesystem.StreamingShell, sink outputSink, toolName, desc string) (tool.BaseTool, error) { + return utils.InferStreamTool(toolName, desc, func(ctx context.Context, input executeManagedArgs) (*schema.StreamReader[string], error) { + req := &filesystem.ExecuteRequest{Command: input.Command} + w := reserveBashOutput(ctx, mgr, sink) + // RunStream owns the returned stream: it forwards work chunks to this caller + // in real time, and on auto-background caps the stream with a notice while + // draining the rest into the task result. A background launch (or timeout) + // is therefore surfaced inline as a final chunk, not as an error. + return mgr.RunStream(ctx, managedRunInput(ctx, input, w), bashStreamWork(streaming, req, w)) + }) +} diff --git a/adk/middlewares/filesystem/bash_run_test.go b/adk/middlewares/filesystem/bash_run_test.go new file mode 100644 index 000000000..e2925744a --- /dev/null +++ b/adk/middlewares/filesystem/bash_run_test.go @@ -0,0 +1,524 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "context" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func intPtr(v int) *int { return &v } + +// findExecuteTool returns the execute tool from a tool set (which, when a Backend +// is configured, also contains the file tools). +func findExecuteTool(t *testing.T, tools []tool.BaseTool) tool.BaseTool { + t.Helper() + for _, to := range tools { + info, err := to.Info(context.Background()) + require.NoError(t, err) + if info.Name == ToolNameExecute { + return to + } + } + t.Fatalf("execute tool %q not found in tool set", ToolNameExecute) + return nil +} + +func waitAllTasks(t *testing.T, mgr *backgroundtask.Manager) { + t.Helper() + require.Eventually(t, func() bool { + for _, task := range mgr.List() { + if task.Status == backgroundtask.StatusRunning { + return false + } + } + return true + }, time.Second, 10*time.Millisecond) +} + +// With a Backend and OutputDir configured, the managed execute tool writes each +// task's output to a file under that directory, and the file is readable back. +func TestManagedExecuteTool_WritesOutputFile(t *testing.T) { + backend := setupTestBackend() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + _ = mgr.Close(ctx) + }() + + tools, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Backend: backend, + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "the output"}}, + Background: &BackgroundConfig{ + Manager: mgr, + OutputStore: backend, + OutputDir: "/tasks", + }, + }) + require.NoError(t, err) + + _, err = invokeTool(t, findExecuteTool(t, tools), `{"command":"echo hi"}`) + require.NoError(t, err) + + tasks := mgr.List() + require.Len(t, tasks, 1) + path := tasks[0].OutputFile + require.NotEmpty(t, path) + + got, err := backend.Read(context.Background(), &filesystem.ReadRequest{FilePath: path}) + require.NoError(t, err) + assert.Equal(t, "the output", got.Content) +} + +// slowShell is a Shell whose Execute blocks for delay (honoring ctx cancellation) +// before returning out. +type slowShell struct { + delay time.Duration + out string +} + +func (s *slowShell) Execute(ctx context.Context, _ *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) { + select { + case <-time.After(s.delay): + return &filesystem.ExecuteResponse{Output: s.out}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// gatedShell is a Shell whose Execute blocks until release is closed (honoring ctx +// cancellation), then returns out. It lets a test hold a background task in the +// running state deterministically, without relying on wall-clock timing. +type gatedShell struct { + release chan struct{} + out string +} + +func (s *gatedShell) Execute(ctx context.Context, _ *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) { + select { + case <-s.release: + return &filesystem.ExecuteResponse{Output: s.out}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func TestManagedExecuteTool_Foreground(t *testing.T) { + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + _ = mgr.Close(ctx) + }() + + tools, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + require.Len(t, tools, 1) + + result, err := invokeTool(t, tools[0], `{"command":"echo hi"}`) + require.NoError(t, err) + assert.Equal(t, "ok", result) + + // The run is tracked by the Manager and tagged as a bash task. + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + assert.Equal(t, "echo hi", tasks[0].Description) + assert.Equal(t, ExecuteTaskType, tasks[0].Type) +} + +func TestManagedExecuteTool_Background(t *testing.T) { + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + backend := setupTestBackend() // so a background launch reports an output path + // A gated shell keeps the background task in the running state until we release + // it, so the launch reliably returns the "running in background" notice rather + // than racing the task to completion. + shell := &gatedShell{release: make(chan struct{}), out: "done"} + tools, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Backend: backend, + Shell: shell, + Background: &BackgroundConfig{ + Manager: mgr, + OutputStore: backend, + OutputDir: "/tasks", + }, + }) + require.NoError(t, err) + + result, err := invokeTool(t, findExecuteTool(t, tools), `{"command":"sleep 1","run_in_background":true}`) + require.NoError(t, err) + assert.Contains(t, result, "running in background") + + // Let the held task finish, then confirm it reaches completion. + close(shell.release) + waitAllTasks(t, mgr) + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.True(t, tasks[0].RunInBackground) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + + // The background-launch message reports the (reserved) output-file path so the + // agent can read it once the task completes. + assert.Contains(t, result, tasks[0].OutputFile) + assert.NotEmpty(t, tasks[0].OutputFile) +} + +// A foreground command that outlives its timeout is moved to the background +// (kept running) when the Manager's ShouldAutoBackground hook permits it. +func TestManagedExecuteTool_TimeoutMovesToBackground(t *testing.T) { + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{ + ForegroundTimeoutMs: intPtr(0), + ShouldAutoBackground: func(context.Context, *backgroundtask.Task) bool { return true }, + }) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + tools, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Shell: &slowShell{delay: 200 * time.Millisecond, out: "slow done"}, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + + // timeout=50ms < 200ms command → moved to background. + result, err := invokeTool(t, tools[0], `{"command":"sleep","timeout":50}`) + require.NoError(t, err) + assert.Contains(t, result, "running in background") + + waitAllTasks(t, mgr) + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + assert.Equal(t, "slow done", tasks[0].Result) +} + +// Without a ShouldAutoBackground hook, a command that outlives its timeout is +// stopped and reported as timed out. +func TestManagedExecuteTool_TimeoutKills(t *testing.T) { + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{ForegroundTimeoutMs: intPtr(0)}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + tools, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Shell: &slowShell{delay: time.Second, out: "never"}, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + + _, err = invokeTool(t, tools[0], `{"command":"sleep","timeout":50}`) + require.Error(t, err) + assert.Contains(t, err.Error(), "timed out") + + waitAllTasks(t, mgr) + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Equal(t, backgroundtask.StatusFailed, tasks[0].Status) +} + +// With a Manager, the execute tool schema gains run_in_background and timeout fields. +// With a StreamingShell backend the managed execute tool is a StreamableTool that +// streams foreground output live while still tracking the run in the Manager. +func TestManagedExecuteTool_StreamingForeground(t *testing.T) { + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + executeTool, err := newManagedExecuteTool(mgr, nil, &mockStreamingShellMultiChunk{}, outputSink{}, "", "") + require.NoError(t, err) + + st, ok := executeTool.(tool.StreamableTool) + require.True(t, ok, "managed execute tool with StreamingShell must be a StreamableTool") + + sr, err := st.StreamableRun(context.Background(), `{"command":"echo hi"}`) + require.NoError(t, err) + got := drainToolStream(t, sr) + assert.Contains(t, got, "chunk1") + assert.Contains(t, got, "chunk3") + + waitAllTasks(t, mgr) + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + assert.Equal(t, ExecuteTaskType, tasks[0].Type) + // The streamed chunks are also the persisted result. + assert.Contains(t, tasks[0].Result, "chunk1") + assert.Contains(t, tasks[0].Result, "chunk3") +} + +// An explicit background launch on a streaming managed tool emits only the +// background notice on the caller's stream; the output lands in the task result. +func TestManagedExecuteTool_StreamingExplicitBackground(t *testing.T) { + backend := setupTestBackend() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + executeTool, err := newManagedExecuteTool(mgr, nil, &mockStreamingShellMultiChunk{}, outputSink{appender: backend, outputDir: "/tasks"}, "", "") + require.NoError(t, err) + st := executeTool.(tool.StreamableTool) + + sr, err := st.StreamableRun(context.Background(), `{"command":"echo hi","run_in_background":true}`) + require.NoError(t, err) + got := drainToolStream(t, sr) + assert.Contains(t, got, "is running in the background") + assert.NotContains(t, got, "moved to the background") + assert.NotContains(t, got, "chunk1") + + waitAllTasks(t, mgr) + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.True(t, tasks[0].RunInBackground) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + assert.Contains(t, tasks[0].Result, "chunk1") + // The streamed output was teed to the output file as it drained in the background. + require.NotEmpty(t, tasks[0].OutputFile) + got2, err := backend.Read(context.Background(), &filesystem.ReadRequest{FilePath: tasks[0].OutputFile}) + require.NoError(t, err) + assert.Contains(t, got2.Content, "chunk1") +} + +// gatedStreamingShell emits "first\n", waits for release, then "second\n" and EOF. +// It lets a test observe interim output: the output file holds a growing prefix +// while the run is mid-stream. +type gatedStreamingShell struct { + release chan struct{} +} + +func (g *gatedStreamingShell) ExecuteStreaming(ctx context.Context, _ *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](2) + go func() { + defer sw.Close() + sw.Send(&filesystem.ExecuteResponse{Output: "first\n"}, nil) + <-g.release + sw.Send(&filesystem.ExecuteResponse{Output: "second\n", ExitCode: ptrOf(0)}, nil) + }() + return sr, nil +} + +// The streaming execute tool tees chunks to the output file as they arrive, so a +// reader sees interim output (a growing prefix) before the run completes. +func TestManagedExecuteTool_StreamingInterimOutput(t *testing.T) { + backend := setupTestBackend() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + gate := &gatedStreamingShell{release: make(chan struct{})} + executeTool, err := newManagedExecuteTool(mgr, nil, gate, outputSink{appender: backend, outputDir: "/tasks"}, "", "") + require.NoError(t, err) + st := executeTool.(tool.StreamableTool) + + sr, err := st.StreamableRun(context.Background(), `{"command":"run"}`) + require.NoError(t, err) + + // Read the first chunk off the caller stream — by then it has also been teed to + // the output file. + first, err := sr.Recv() + require.NoError(t, err) + assert.Contains(t, first, "first") + + tasks := mgr.List() + require.Len(t, tasks, 1) + path := tasks[0].OutputFile + require.NotEmpty(t, path) + + // Interim: the file holds the first chunk but not yet the second. + require.Eventually(t, func() bool { + got, readErr := backend.Read(context.Background(), &filesystem.ReadRequest{FilePath: path}) + return readErr == nil && strings.Contains(got.Content, "first") + }, time.Second, 5*time.Millisecond) + interim, err := backend.Read(context.Background(), &filesystem.ReadRequest{FilePath: path}) + require.NoError(t, err) + assert.NotContains(t, interim.Content, "second", "second chunk must not be present before release") + + // Release the rest and drain. + close(gate.release) + for { + if _, err := sr.Recv(); err == io.EOF { + break + } else { + require.NoError(t, err) + } + } + + waitAllTasks(t, mgr) + final, err := backend.Read(context.Background(), &filesystem.ReadRequest{FilePath: path}) + require.NoError(t, err) + assert.Contains(t, final.Content, "first") + assert.Contains(t, final.Content, "second") +} + +// drainToolStream reads a tool's string stream to EOF and returns the joined text. +func drainToolStream(t *testing.T, sr *schema.StreamReader[string]) string { + t.Helper() + defer sr.Close() + var b strings.Builder + for { + chunk, err := sr.Recv() + if err == io.EOF { + return b.String() + } + require.NoError(t, err) + b.WriteString(chunk) + } +} + +func TestManagedExecuteTool_Schema(t *testing.T) { + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + _ = mgr.Close(ctx) + }() + + executeTool, err := newManagedExecuteTool(mgr, &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, nil, outputSink{}, "", "") + require.NoError(t, err) + + info, err := executeTool.Info(context.Background()) + require.NoError(t, err) + js, err := info.ParamsOneOf.ToJSONSchema() + require.NoError(t, err) + assert.Equal(t, 3, js.Properties.Len()) + _, ok := js.Properties.Get("command") + assert.True(t, ok) + _, ok = js.Properties.Get("run_in_background") + assert.True(t, ok) + _, ok = js.Properties.Get("timeout") + assert.True(t, ok) +} + +// Without a Manager, the execute tool is command-only and untracked. +func TestExecuteTool_NoManager_NotTracked(t *testing.T) { + tools, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + }) + require.NoError(t, err) + require.Len(t, tools, 1) + + result, err := invokeTool(t, tools[0], `{"command":"echo hi"}`) + require.NoError(t, err) + assert.Equal(t, "ok", result) +} + +// failingAppender wraps a Backend but fails Append after failAfter successful +// appends (failAfter=0 fails the very first append, i.e. the reservation write). +// Reads delegate to the backend so the partial file is still observable. +type failingAppender struct { + backend *filesystem.InMemoryBackend + failAfter int + calls int +} + +func (f *failingAppender) Append(ctx context.Context, req *filesystem.AppendRequest) error { + if f.calls >= f.failAfter { + f.calls++ + return errors.New("append failed") + } + f.calls++ + return f.backend.Append(ctx, req) +} + +// When the up-front reservation write fails, the task advertises no output file, +// so consumers fall back to the in-memory Result. +func TestManagedExecuteTool_ReservationFailure_NoOutputFile(t *testing.T) { + backend := setupTestBackend() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + appender := &failingAppender{backend: backend, failAfter: 0} + executeTool, err := newManagedExecuteTool(mgr, &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "the output"}}, nil, + outputSink{appender: appender, outputDir: "/tasks"}, "", "") + require.NoError(t, err) + + result, err := invokeTool(t, executeTool, `{"command":"echo hi"}`) + require.NoError(t, err) + assert.Equal(t, "the output", result) + + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Empty(t, tasks[0].OutputFile, "reservation failure must leave OutputFile unset") + assert.Empty(t, tasks[0].OutputFileErr) + assert.Equal(t, "the output", tasks[0].Result) +} + +// When a write to the output file fails after reservation, the file is marked +// unreliable (OutputFileErr set) while the in-memory Result stays complete. +func TestManagedExecuteTool_WriteFailure_MarksUnreliable(t *testing.T) { + backend := setupTestBackend() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(ctx) + }() + + // failAfter=1: the reservation write succeeds, the result write fails. + appender := &failingAppender{backend: backend, failAfter: 1} + executeTool, err := newManagedExecuteTool(mgr, &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "the output"}}, nil, + outputSink{appender: appender, outputDir: "/tasks"}, "", "") + require.NoError(t, err) + + result, err := invokeTool(t, executeTool, `{"command":"echo hi"}`) + require.NoError(t, err) + assert.Equal(t, "the output", result) + + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.NotEmpty(t, tasks[0].OutputFile, "the path was reserved, so it is still recorded") + assert.NotEmpty(t, tasks[0].OutputFileErr, "the failed write must mark the file unreliable") + assert.Equal(t, "the output", tasks[0].Result, "Result stays complete regardless of file writes") +} diff --git a/adk/middlewares/filesystem/filesystem.go b/adk/middlewares/filesystem/filesystem.go index b9d64ab24..48bd4b691 100644 --- a/adk/middlewares/filesystem/filesystem.go +++ b/adk/middlewares/filesystem/filesystem.go @@ -29,6 +29,7 @@ import ( "strings" "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/backgroundtask" "github.com/cloudwego/eino/adk/filesystem" "github.com/cloudwego/eino/adk/internal" "github.com/cloudwego/eino/components/tool" @@ -72,6 +73,42 @@ type ToolConfig struct { Disable bool } +// ExecuteToolConfig configures the execute tool. +// +// The execute tool's input schema is determined by whether a background-task +// Manager is configured on the middleware: without a Manager the tool accepts +// only a command; with a Manager it additionally accepts a run_in_background +// flag and routes runs through the Manager for lifecycle tracking. +type ExecuteToolConfig struct { + ToolConfig +} + +// BackgroundConfig enables background-task execution for the execute tool. +// +// When set, the execute tool gains a run_in_background field and routes runs +// through the shared Manager, so background and auto-background runs are tracked +// and visible to the task_output/task_stop control tools. With a StreamingShell +// backend the foreground phase still streams in real time; once a run moves to the +// background its stream is capped with a notice and the rest is collected into the +// task result. +type BackgroundConfig struct { + // Manager is the shared background-task Manager. Required (a nil Manager is the + // same as no BackgroundConfig). It may be shared with other middlewares (e.g. + // subagent) for a unified task-ID space; wire the backgroundtask control + // middleware once, bound to the same Manager. + Manager *backgroundtask.Manager + + // OutputStore and OutputDir, when both set, give every managed run an output + // file at OutputDir/.output: streaming runs append their chunks to it as + // they arrive (interim output), buffered runs append their result on completion. + // The path is recorded on Task.OutputFile and surfaced in the background notice. + // OutputStore is a filesystem.Appender (filesystem.InMemoryBackend implements + // it); supply your own to direct output elsewhere. When either is unset, runs + // have no output file. + OutputStore filesystem.Appender + OutputDir string +} + // Config is the configuration for the filesystem middleware type Config struct { // Backend provides filesystem operations used by tools and offloading. @@ -90,6 +127,11 @@ type Config struct { // Mutually exclusive with Shell. StreamingShell filesystem.StreamingShell + // Background configures background-task execution for the execute tool. When + // nil, execute runs only foreground (blocking) and is not tracked. See + // BackgroundConfig. + Background *BackgroundConfig + // LsToolConfig configures the ls tool // optional LsToolConfig *ToolConfig @@ -110,6 +152,9 @@ type Config struct { // GrepToolConfig configures the grep tool // optional GrepToolConfig *ToolConfig + // ExecuteToolConfig configures the execute tool + // optional + ExecuteToolConfig *ExecuteToolConfig // WithoutLargeToolResultOffloading disables automatic offloading of large tool result to Backend // optional, false(enabled) by default @@ -155,8 +200,8 @@ func (c *Config) Validate() error { if c == nil { return errors.New("config should not be nil") } - if c.Backend == nil { - return errors.New("backend should not be nil") + if c.Backend == nil && c.Shell == nil && c.StreamingShell == nil { + return errors.New("at least one of backend, shell, or streaming shell should be set") } if c.StreamingShell != nil && c.Shell != nil { return errors.New("shell and streaming shell should not be both set") @@ -179,12 +224,14 @@ func NewMiddleware(ctx context.Context, config *Config) (adk.AgentMiddleware, er Backend: config.Backend, Shell: config.Shell, StreamingShell: config.StreamingShell, + Background: config.Background, LsToolConfig: config.LsToolConfig, ReadFileToolConfig: config.ReadFileToolConfig, WriteFileToolConfig: config.WriteFileToolConfig, EditFileToolConfig: config.EditFileToolConfig, GlobToolConfig: config.GlobToolConfig, GrepToolConfig: config.GrepToolConfig, + ExecuteToolConfig: config.ExecuteToolConfig, CustomSystemPrompt: config.CustomSystemPrompt, CustomLsToolDesc: config.CustomLsToolDesc, CustomReadFileToolDesc: config.CustomReadFileToolDesc, @@ -207,7 +254,7 @@ func NewMiddleware(ctx context.Context, config *Config) (adk.AgentMiddleware, er AdditionalTools: ts, } - if !config.WithoutLargeToolResultOffloading { + if config.Backend != nil && !config.WithoutLargeToolResultOffloading { m.WrapToolCall = newToolResultOffloading(ctx, &toolResultOffloadingConfig{ Backend: config.Backend, TokenLimit: config.LargeToolResultOffloadingTokenLimit, @@ -221,18 +268,25 @@ func NewMiddleware(ctx context.Context, config *Config) (adk.AgentMiddleware, er // MiddlewareConfig is the configuration for the filesystem middleware type MiddlewareConfig struct { // Backend provides filesystem operations used by tools and offloading. - // required + // At least one of Backend, Shell, or StreamingShell must be set. Backend filesystem.Backend // Shell provides shell command execution capability. // If set, an execute tool will be registered to support shell command execution. - // optional, mutually exclusive with StreamingShell + // At least one of Backend, Shell, or StreamingShell must be set. + // Mutually exclusive with StreamingShell. Shell filesystem.Shell // StreamingShell provides streaming shell command execution capability. // If set, a streaming execute tool will be registered for real-time output. - // optional, mutually exclusive with Shell + // At least one of Backend, Shell, or StreamingShell must be set. + // Mutually exclusive with Shell. StreamingShell filesystem.StreamingShell + // Background configures background-task execution for the execute tool. When + // nil, execute runs only foreground (blocking) and is not tracked. See + // BackgroundConfig. + Background *BackgroundConfig + // LsToolConfig configures the ls tool // optional LsToolConfig *ToolConfig @@ -253,6 +307,9 @@ type MiddlewareConfig struct { // GrepToolConfig configures the grep tool // optional GrepToolConfig *ToolConfig + // ExecuteToolConfig configures the execute tool + // optional + ExecuteToolConfig *ExecuteToolConfig // UseMultiModalRead enables multimodal read_file tool (EnhancedInvokableTool). // When true, read_file returns results via schema.ToolResult.Parts instead of plain text string. @@ -306,8 +363,8 @@ func (c *MiddlewareConfig) Validate() error { if c == nil { return errors.New("config should not be nil") } - if c.Backend == nil { - return errors.New("backend should not be nil") + if c.Backend == nil && c.Shell == nil && c.StreamingShell == nil { + return errors.New("at least one of backend, shell, or streaming shell should be set") } if c.StreamingShell != nil && c.Shell != nil { return errors.New("shell and streaming shell should not be both set") @@ -350,7 +407,7 @@ func (c *MiddlewareConfig) mergeToolConfigWithDesc( // - More flexible extension points compared to the struct-based AgentMiddleware // // The middleware provides filesystem tools (ls, read_file, write_file, edit_file, glob, grep) -// and optionally an execute tool if the Backend implements ShellBackend or StreamingShellBackend. +// when Backend is set, and an execute tool when Shell or StreamingShell is set. func NewTyped[M adk.MessageType](ctx context.Context, config *MiddlewareConfig) (adk.TypedChatModelAgentMiddleware[M], error) { err := config.Validate() if err != nil { @@ -381,7 +438,7 @@ func NewTyped[M adk.MessageType](ctx context.Context, config *MiddlewareConfig) // - More flexible extension points compared to the struct-based AgentMiddleware // // The middleware provides filesystem tools (ls, read_file, write_file, edit_file, glob, grep) -// and optionally an execute tool if the Backend implements ShellBackend or StreamingShellBackend. +// when Backend is set, and an execute tool when Shell or StreamingShell is set. // // Example usage: // @@ -402,7 +459,7 @@ type typedFilesystemMiddleware[M adk.MessageType] struct { additionalTools []tool.BaseTool } -func (m *typedFilesystemMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { +func (m *typedFilesystemMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { if runCtx == nil { return ctx, runCtx, nil } @@ -415,9 +472,8 @@ func (m *typedFilesystemMiddleware[M]) BeforeAgent(ctx context.Context, runCtx * return ctx, &nRunCtx, nil } -// toolSpec defines a specification for creating a filesystem tool. -// It unifies the tool creation process by encapsulating the tool configuration, -// legacy descriptor, and the creation function. +// toolSpec describes how to construct one filesystem tool, including its +// configuration, legacy descriptor, and constructor. type toolSpec struct { config *ToolConfig legacyDesc *string @@ -503,37 +559,61 @@ func getFilesystemTools(_ context.Context, middlewareConfig *MiddlewareConfig) ( } } - // Create execute tool if Shell or StreamingShell is available - if middlewareConfig.StreamingShell != nil { - executeDesc, err := selectToolDesc("", ExecuteToolDesc, ExecuteToolDescChinese) + if middlewareConfig.StreamingShell != nil || middlewareConfig.Shell != nil { + executeTool, err := createExecuteTool(middlewareConfig) if err != nil { return nil, err } - - executeTool, err := newStreamingExecuteTool(middlewareConfig.StreamingShell, ToolNameExecute, executeDesc) - if err != nil { - return nil, err + if executeTool != nil { + tools = append(tools, executeTool) } - tools = append(tools, executeTool) - } else if middlewareConfig.Shell != nil { - executeDesc, err := selectToolDesc("", ExecuteToolDesc, ExecuteToolDescChinese) - if err != nil { - return nil, err + } + + return tools, nil +} + +func createExecuteTool(middlewareConfig *MiddlewareConfig) (tool.BaseTool, error) { + executeConfig := middlewareConfig.ExecuteToolConfig + if executeConfig == nil { + executeConfig = &ExecuteToolConfig{} + } + if executeConfig.Disable { + return nil, nil + } + return getOrCreateTool(executeConfig.CustomTool, func() (tool.BaseTool, error) { + desc := "" + if executeConfig.Desc != nil { + desc = *executeConfig.Desc } - executeTool, err := newExecuteTool(middlewareConfig.Shell, ToolNameExecute, executeDesc) - if err != nil { - return nil, err + // When a shared Manager is configured, the execute tool exposes a + // run_in_background field and routes runs through the Manager, so + // background/auto-background runs are tracked and visible to the + // task_output/task_stop control tools. Without a Manager the tool is + // command-only with no background support. + if middlewareConfig.Background != nil && middlewareConfig.Background.Manager != nil { + return newManagedExecuteTool( + middlewareConfig.Background.Manager, + middlewareConfig.Shell, + middlewareConfig.StreamingShell, + outputSink{ + appender: middlewareConfig.Background.OutputStore, + outputDir: middlewareConfig.Background.OutputDir, + }, + executeConfig.Name, + desc, + ) } - tools = append(tools, executeTool) - } - return tools, nil + if middlewareConfig.StreamingShell != nil { + return newStreamingExecuteTool(middlewareConfig.StreamingShell, executeConfig.Name, desc) + } + return newExecuteTool(middlewareConfig.Shell, executeConfig.Name, desc) + }) } -// createToolFromSpec creates a tool instance based on the provided toolSpec. -// It handles configuration merging (ToolConfig + legacy Desc), checks if the tool -// is disabled, and prioritizes CustomTool over the default implementation. +// createToolFromSpec creates a tool from spec, applying configuration merging, +// disable handling, and CustomTool precedence. func createToolFromSpec(middlewareConfig *MiddlewareConfig, spec toolSpec) (tool.BaseTool, error) { mergedConfig := middlewareConfig.mergeToolConfigWithDesc(spec.config, spec.legacyDesc) @@ -997,7 +1077,19 @@ func newGrepTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool } type executeArgs struct { - Command string `json:"command"` + Command string `json:"command" jsonschema:"required" jsonschema_description:"The command to execute"` +} + +// executeManagedArgs is the execute tool input used when a background-task +// Manager is configured: the model may additionally request background execution. +type executeManagedArgs struct { + executeArgs + RunInBackground bool `json:"run_in_background,omitempty" jsonschema_description:"Set to true to run the command in the background. You will be notified when it completes; use task_output to query it and task_stop to cancel it."` + // TimeoutMS is the foreground budget in milliseconds. When omitted, the configured + // default applies. Ignored when run_in_background is true. What happens at the + // deadline (move to background vs. stop) is decided by the Manager's + // ShouldAutoBackground policy and is intentionally not surfaced to the model. + TimeoutMS int `json:"timeout,omitempty" jsonschema_description:"Optional timeout in milliseconds. The maximum time to wait for the command. Omit to use the default."` } func newExecuteTool(sb filesystem.Shell, name string, desc string) (tool.BaseTool, error) { @@ -1007,13 +1099,10 @@ func newExecuteTool(sb filesystem.Shell, name string, desc string) (tool.BaseToo return nil, err } return utils.InferTool(toolName, d, func(ctx context.Context, input executeArgs) (string, error) { - result, err := sb.Execute(ctx, &filesystem.ExecuteRequest{ - Command: input.Command, - }) + result, err := sb.Execute(ctx, &filesystem.ExecuteRequest{Command: input.Command}) if err != nil { return "", err } - return convExecuteResponse(result), nil }) } @@ -1024,10 +1113,23 @@ func newStreamingExecuteTool(sb filesystem.StreamingShell, name string, desc str if err != nil { return nil, err } - return utils.InferStreamTool(toolName, d, func(ctx context.Context, input executeArgs) (*schema.StreamReader[string], error) { - result, err := sb.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{ - Command: input.Command, - }) + return newStreamingExecuteToolWithRun(sb, toolName, d, func(input executeArgs) (*filesystem.ExecuteRequest, error) { + return &filesystem.ExecuteRequest{Command: input.Command}, nil + }) +} + +func newStreamingExecuteToolWithRun[T any]( + sb filesystem.StreamingShell, + toolName string, + desc string, + newRequest func(input T) (*filesystem.ExecuteRequest, error), +) (tool.BaseTool, error) { + return utils.InferStreamTool(toolName, desc, func(ctx context.Context, input T) (*schema.StreamReader[string], error) { + req, err := newRequest(input) + if err != nil { + return nil, err + } + result, err := sb.ExecuteStreaming(ctx, req) if err != nil { return nil, err } @@ -1061,23 +1163,14 @@ func newStreamingExecuteTool(sb filesystem.StreamingShell, name string, desc str exitCode = chunk.ExitCode } - parts := make([]string, 0, 2) - if chunk.Output != "" { - parts = append(parts, chunk.Output) - } - if chunk.Truncated { - parts = append(parts, "[Output was truncated due to size limits]") - } - if len(parts) > 0 { - sw.Send(strings.Join(parts, "\n"), nil) + if text := formatExecChunk(chunk.Output, chunk.Truncated); text != "" { + sw.Send(text, nil) hasSentContent = true } } - if exitCode != nil && *exitCode != 0 { - sw.Send(fmt.Sprintf("\n[Command failed with exit code %d]", *exitCode), nil) - } else if !hasSentContent { - sw.Send("[Command executed successfully with no output]", nil) + if note := execTerminalNote(exitCode, hasSentContent); note != "" { + sw.Send(note, nil) } }() @@ -1085,21 +1178,55 @@ func newStreamingExecuteTool(sb filesystem.StreamingShell, name string, desc str }) } +// Markers appended to execute-tool output. Shared by the buffered (convExecuteResponse), +// streaming (newStreamingExecuteToolWithRun), and managed (bashStreamWork) paths. +const ( + outputTruncatedNote = "[Output was truncated due to size limits]" + commandFailedFmt = "[Command failed with exit code %d]" + noCommandOutputNote = "[Command executed successfully with no output]" +) + +// formatExecChunk renders one streamed ExecuteResponse chunk to the text to emit, +// or "" when the chunk carries nothing. +func formatExecChunk(output string, truncated bool) string { + parts := make([]string, 0, 2) + if output != "" { + parts = append(parts, output) + } + if truncated { + parts = append(parts, outputTruncatedNote) + } + return strings.Join(parts, "\n") +} + +// execTerminalNote returns the trailing text for a finished command: a failure note +// for a non-zero exit code, the no-output message when nothing was emitted, or "" +// otherwise. +func execTerminalNote(exitCode *int, hasContent bool) string { + if exitCode != nil && *exitCode != 0 { + return "\n" + fmt.Sprintf(commandFailedFmt, *exitCode) + } + if !hasContent { + return noCommandOutputNote + } + return "" +} + func convExecuteResponse(response *filesystem.ExecuteResponse) string { if response == nil { return "" } parts := []string{response.Output} if response.ExitCode != nil && *response.ExitCode != 0 { - parts = append(parts, fmt.Sprintf("[Command failed with exit code %d]", *response.ExitCode)) + parts = append(parts, fmt.Sprintf(commandFailedFmt, *response.ExitCode)) } if response.Truncated { - parts = append(parts, "[Output was truncated due to size limits]") + parts = append(parts, outputTruncatedNote) } result := strings.Join(parts, "\n") if result == "" && (response.ExitCode == nil || *response.ExitCode == 0) { - return "[Command executed successfully with no output]" + return noCommandOutputNote } return result } diff --git a/adk/middlewares/filesystem/filesystem_test.go b/adk/middlewares/filesystem/filesystem_test.go index cb59353ca..b70efc0ca 100644 --- a/adk/middlewares/filesystem/filesystem_test.go +++ b/adk/middlewares/filesystem/filesystem_test.go @@ -577,6 +577,37 @@ func TestExecuteTool(t *testing.T) { } } +func TestExecuteToolSchema_NoManager(t *testing.T) { + ctx := context.Background() + + t.Run("schema is command only", func(t *testing.T) { + executeTool, err := newExecuteTool(&mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, "", "") + assert.NoError(t, err) + + info, err := executeTool.Info(ctx) + assert.NoError(t, err) + js, err := info.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + assert.NotNil(t, js) + assert.Equal(t, 1, js.Properties.Len()) + _, ok := js.Properties.Get("command") + assert.True(t, ok) + _, ok = js.Properties.Get("run_in_background") + assert.False(t, ok) + }) + + t.Run("forwards only the command to the backend", func(t *testing.T) { + shell := &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}} + executeTool, err := newExecuteTool(shell, "", "") + assert.NoError(t, err) + + result, err := invokeTool(t, executeTool, `{"command": "echo ok"}`) + assert.NoError(t, err) + assert.Equal(t, "ok", result) + assert.Equal(t, "echo ok", shell.req.Command) + }) +} + func ptrOf[T any](t T) *T { return &t } @@ -584,9 +615,11 @@ func ptrOf[T any](t T) *T { type mockShellBackend struct { filesystem.Backend resp *filesystem.ExecuteResponse + req *filesystem.ExecuteRequest } func (m *mockShellBackend) Execute(ctx context.Context, req *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) { + m.req = req return m.resp, nil } @@ -656,6 +689,101 @@ func TestGetFilesystemTools(t *testing.T) { }) } +func TestExecuteToolConfig(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("disable skips execute registration", func(t *testing.T) { + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Backend: backend, + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + ExecuteToolConfig: &ExecuteToolConfig{ + ToolConfig: ToolConfig{Disable: true}, + }, + }) + assert.NoError(t, err) + assert.Len(t, tools, 6) + for _, to := range tools { + info, err := to.Info(ctx) + assert.NoError(t, err) + assert.NotEqual(t, ToolNameExecute, info.Name) + } + }) + + t.Run("custom tool overrides built-in execute", func(t *testing.T) { + customTool, err := newLsTool(backend, "custom_execute", "custom execute") + assert.NoError(t, err) + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + ExecuteToolConfig: &ExecuteToolConfig{ + ToolConfig: ToolConfig{CustomTool: customTool}, + }, + }) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, customTool, tools[0]) + }) + + t.Run("name and desc apply to built-in execute", func(t *testing.T) { + desc := "custom execute desc" + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + ExecuteToolConfig: &ExecuteToolConfig{ + ToolConfig: ToolConfig{ + Name: "run", + Desc: &desc, + }, + }, + }) + assert.NoError(t, err) + assert.Len(t, tools, 1) + info, err := tools[0].Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "run", info.Name) + assert.Equal(t, desc, info.Desc) + }) + + t.Run("deprecated config passes execute tool config through", func(t *testing.T) { + m, err := NewMiddleware(ctx, &Config{ + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + ExecuteToolConfig: &ExecuteToolConfig{ + ToolConfig: ToolConfig{Name: "run"}, + }, + }) + assert.NoError(t, err) + assert.Len(t, m.AdditionalTools, 1) + info, err := m.AdditionalTools[0].Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "run", info.Name) + js, err := info.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + _, ok := js.Properties.Get("command") + assert.True(t, ok) + }) +} + +func TestGetFilesystemTools_NoExecuteLifecycleTools(t *testing.T) { + ctx := context.Background() + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Backend: setupTestBackend(), + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + }) + assert.NoError(t, err) + + toolNames := make(map[string]bool) + for _, to := range tools { + info, err := to.Info(ctx) + assert.NoError(t, err) + toolNames[info.Name] = true + } + + assert.True(t, toolNames[ToolNameExecute]) + assert.False(t, toolNames["execute_output"]) + assert.False(t, toolNames["execute_wait"]) + assert.False(t, toolNames["execute_stop"]) + assert.False(t, toolNames["execute_list"]) +} + func TestNew(t *testing.T) { ctx := context.Background() backend := setupTestBackend() @@ -666,10 +794,24 @@ func TestNew(t *testing.T) { assert.Contains(t, err.Error(), "config should not be nil") }) - t.Run("nil backend returns error", func(t *testing.T) { + t.Run("all execution backends nil returns error", func(t *testing.T) { _, err := New(ctx, &MiddlewareConfig{Backend: nil}) assert.Error(t, err) - assert.Contains(t, err.Error(), "backend should not be nil") + assert.Contains(t, err.Error(), "at least one of backend, shell, or streaming shell should be set") + }) + + t.Run("shell-only config registers execute tool", func(t *testing.T) { + m, err := New(ctx, &MiddlewareConfig{ + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + }) + assert.NoError(t, err) + + fm, ok := m.(*typedFilesystemMiddleware[*schema.Message]) + assert.True(t, ok) + assert.Len(t, fm.additionalTools, 1) + info, err := fm.additionalTools[0].Info(ctx) + assert.NoError(t, err) + assert.Equal(t, ToolNameExecute, info.Name) }) t.Run("valid config with default settings", func(t *testing.T) { @@ -717,7 +859,7 @@ func TestFilesystemMiddleware_BeforeAgent(t *testing.T) { m, err := New(ctx, &MiddlewareConfig{Backend: backend}) assert.NoError(t, err) - runCtx := &adk.ChatModelAgentContext{ + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ Instruction: "Original instruction", Tools: nil, } @@ -1637,7 +1779,7 @@ func TestGetFilesystemTools_NilBackend(t *testing.T) { Backend: nil, StreamingShell: mockSS, } - // Validate should fail, but getFilesystemTools itself handles nil backend gracefully + assert.NoError(t, config.Validate()) tools, err := getFilesystemTools(ctx, config) assert.NoError(t, err) // Only execute tool should be returned since backend is nil @@ -1700,9 +1842,12 @@ func TestGetFilesystemTools_PartialDisable(t *testing.T) { assert.Contains(t, toolNames, ToolNameGrep) } -type mockStreamingShell struct{} +type mockStreamingShell struct { + req *filesystem.ExecuteRequest +} func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + m.req = input sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) go func() { defer sw.Close() @@ -1946,6 +2091,25 @@ func TestNewStreamingExecuteTool(t *testing.T) { assert.Equal(t, "custom_execute", info.Name) assert.Equal(t, "custom desc", info.Desc) }) + + t.Run("streaming forwards only command", func(t *testing.T) { + streamingShell := &mockStreamingShell{} + executeTool, err := newStreamingExecuteTool(streamingShell, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "echo hello"}`) + assert.NoError(t, err) + defer sr.Close() + for { + _, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + } + assert.Equal(t, "echo hello", streamingShell.req.Command) + }) } func TestNew_StreamingShell(t *testing.T) { @@ -1984,10 +2148,10 @@ func TestNewMiddleware_Validation(t *testing.T) { assert.Contains(t, err.Error(), "config should not be nil") }) - t.Run("nil backend returns error", func(t *testing.T) { + t.Run("all execution backends nil returns error", func(t *testing.T) { _, err := NewMiddleware(ctx, &Config{Backend: nil}) assert.Error(t, err) - assert.Contains(t, err.Error(), "backend should not be nil") + assert.Contains(t, err.Error(), "at least one of backend, shell, or streaming shell should be set") }) t.Run("both Shell and StreamingShell returns error", func(t *testing.T) { @@ -2010,11 +2174,11 @@ func TestMiddlewareConfig_Validate(t *testing.T) { assert.Contains(t, err.Error(), "config should not be nil") }) - t.Run("nil backend returns error", func(t *testing.T) { + t.Run("all execution backends nil returns error", func(t *testing.T) { c := &MiddlewareConfig{} err := c.Validate() assert.Error(t, err) - assert.Contains(t, err.Error(), "backend should not be nil") + assert.Contains(t, err.Error(), "at least one of backend, shell, or streaming shell should be set") }) t.Run("both shells returns error", func(t *testing.T) { @@ -2035,6 +2199,14 @@ func TestMiddlewareConfig_Validate(t *testing.T) { err := c.Validate() assert.NoError(t, err) }) + + t.Run("shell-only config passes", func(t *testing.T) { + c := &MiddlewareConfig{ + Shell: &mockShellBackend{}, + } + err := c.Validate() + assert.NoError(t, err) + }) } func TestNewStreamingExecuteTool_MultipleChunks(t *testing.T) { @@ -2134,11 +2306,11 @@ func TestConfig_Validate(t *testing.T) { assert.Error(t, err) }) - t.Run("nil backend returns error", func(t *testing.T) { + t.Run("all execution backends nil returns error", func(t *testing.T) { c := &Config{} err := c.Validate() assert.Error(t, err) - assert.Contains(t, err.Error(), "backend should not be nil") + assert.Contains(t, err.Error(), "at least one of backend, shell, or streaming shell should be set") }) t.Run("both shells returns error", func(t *testing.T) { @@ -2158,6 +2330,14 @@ func TestConfig_Validate(t *testing.T) { err := c.Validate() assert.NoError(t, err) }) + + t.Run("shell-only config passes", func(t *testing.T) { + c := &Config{ + Shell: &mockShellBackend{}, + } + err := c.Validate() + assert.NoError(t, err) + }) } func TestGetFilesystemTools_CustomToolWithShell(t *testing.T) { @@ -2256,6 +2436,30 @@ func TestNewMiddleware_WithShell(t *testing.T) { assert.NoError(t, err) assert.Len(t, m.AdditionalTools, 7) }) + + t.Run("shell-only config skips large tool result offloading", func(t *testing.T) { + m, err := NewMiddleware(ctx, &Config{ + Shell: &mockShellBackend{resp: &filesystem.ExecuteResponse{Output: "ok"}}, + }) + assert.NoError(t, err) + assert.Len(t, m.AdditionalTools, 1) + assert.Nil(t, m.WrapToolCall.Invokable) + assert.Nil(t, m.WrapToolCall.Streamable) + assert.Nil(t, m.WrapToolCall.EnhancedInvokable) + assert.Nil(t, m.WrapToolCall.EnhancedStreamable) + }) + + t.Run("streaming shell-only config skips large tool result offloading", func(t *testing.T) { + m, err := NewMiddleware(ctx, &Config{ + StreamingShell: &mockStreamingShell{}, + }) + assert.NoError(t, err) + assert.Len(t, m.AdditionalTools, 1) + assert.Nil(t, m.WrapToolCall.Invokable) + assert.Nil(t, m.WrapToolCall.Streamable) + assert.Nil(t, m.WrapToolCall.EnhancedInvokable) + assert.Nil(t, m.WrapToolCall.EnhancedStreamable) + }) } func TestNewExecuteTool_ShellError(t *testing.T) { diff --git a/adk/middlewares/filesystem/prompt.go b/adk/middlewares/filesystem/prompt.go index a20d6d7d8..244013b48 100644 --- a/adk/middlewares/filesystem/prompt.go +++ b/adk/middlewares/filesystem/prompt.go @@ -261,6 +261,94 @@ Bad examples (avoid these): - execute(command="python /path/to/script.py") - execute(command="npm install && npm test") +不好的示例（避免这些）： +- execute(command="cd /foo/bar && pytest tests") # 改用绝对路径 +- execute(command="cat file.txt") # 改用 read_file 工具 +- execute(command="find . -name '*.py'") # 改用 glob 工具 +- execute(command="grep -r 'pattern' .") # 改用 grep 工具 +` + + ManagedExecuteToolDesc = ` +Executes a given command in the sandbox environment with proper handling and security measures. + +Before executing the command, please follow these steps: + +1. Directory Verification: +- If the command will create new directories or files, first use the ls tool to verify the parent directory exists and is the correct location +- For example, before running "mkdir foo/bar", first use ls to check that "foo" exists and is the intended parent directory + +2. Command Execution: +- Always quote file paths that contain spaces with double quotes (e.g., cd "path with spaces/file.txt") +- Examples of proper quoting: +- cd "/Users/name/My Documents" (correct) +- cd /Users/name/My Documents (incorrect - will fail) +- python "/path/with spaces/script.py" (correct) +- python /path/with spaces/script.py (incorrect - will fail) +- After ensuring proper quoting, execute the command +- Capture the output of the command + +Usage notes: +- The command parameter is required +- Set run_in_background=true for servers, watchers, and other long-running commands you do not need to wait for. You will be notified when it completes; use the task_output tool to check its status or retrieve its result, and the task_stop tool to cancel it. +- The optional timeout parameter (in milliseconds) sets the maximum time to wait for the command. Omit to use the default. +- Commands run in an isolated sandbox environment +- Returns combined stdout/stderr output with exit code +- If the output is very large, it may be truncated +- VERY IMPORTANT: You MUST avoid using search commands like find and grep. Instead use the grep, glob tools to search. You MUST avoid read tools like cat, head, tail, and use read_file to read files. +- When issuing multiple commands, use the ';' or '&&' operator to separate them. DO NOT use newlines (newlines are ok in quoted strings) +- Use '&&' when commands depend on each other (e.g., "mkdir dir && cd dir") +- Use ';' only when you need to run commands sequentially but don't care if earlier commands fail +- Try to maintain your current working directory throughout the session by using absolute paths and avoiding usage of cd + +Examples: +Good examples: +- execute(command="pytest /foo/bar/tests") +- execute(command="npm run dev", run_in_background=true) + +Bad examples (avoid these): +- execute(command="cd /foo/bar && pytest tests") # Use absolute path instead +- execute(command="cat file.txt") # Use read_file tool instead +- execute(command="find . -name '*.py'") # Use glob tool instead +- execute(command="grep -r 'pattern' .") # Use grep tool instead +` + + ManagedExecuteToolDescChinese = ` +在沙箱环境中执行给定命令，具有适当的处理和安全措施。 + +执行命令前，请按照以下步骤操作： + +1. 目录验证： +- 如果命令将创建新目录或文件，首先使用 ls 工具验证父目录是否存在且是正确的位置 +- 例如，在运行 "mkdir foo/bar" 之前，首先使用 ls 检查 "foo" 是否存在且是预期的父目录 + +2. 命令执行： +- 始终用双引号引用包含空格的文件路径（例如，cd "path with spaces/file.txt"） +- 正确引用的示例： +- cd "/Users/name/My Documents"（正确） +- cd /Users/name/My Documents（错误 - 将失败） +- python "/path/with spaces/script.py"（正确） +- python /path/with spaces/script.py（错误 - 将失败） +- 确保正确引用后，执行命令 +- 捕获命令的输出 + +使用说明： +- command 参数是必需的 +- 对于服务器、监听器等你无需等待的长时间运行命令，设置 run_in_background=true。命令完成时你会收到通知；使用 task_output 工具查询其状态或获取结果，使用 task_stop 工具取消它。 +- 可选的 timeout 参数（毫秒）设置等待命令的最长时间。不传则使用默认值。 +- 命令在隔离的沙箱环境中运行 +- 返回合并的 stdout/stderr 输出和退出代码 +- 如果输出非常大，可能会被截断 +- 非常重要：你必须避免使用 find 和 grep 等搜索命令。请改用 grep、glob 工具进行搜索。你必须避免使用 cat、head、tail 等读取工具，请使用 read_file 读取文件 +- 发出多个命令时，使用 ';' 或 '&&' 运算符分隔它们。不要使用换行符（引号字符串中的换行符是可以的） +- 当命令相互依赖时使用 '&&'（例如，"mkdir dir && cd dir"） +- 仅当你需要按顺序运行命令但不关心早期命令是否失败时使用 ';' +- 尝试通过使用绝对路径并避免使用 cd 来在整个会话中保持当前工作目录 + +示例： +好的示例： +- execute(command="pytest /foo/bar/tests") +- execute(command="npm run dev", run_in_background=true) + 不好的示例（避免这些）： - execute(command="cd /foo/bar && pytest tests") # 改用绝对路径 - execute(command="cat file.txt") # 改用 read_file 工具 diff --git a/adk/middlewares/modeltimeout/modeltimeout.go b/adk/middlewares/modeltimeout/modeltimeout.go new file mode 100644 index 000000000..2501d2d63 --- /dev/null +++ b/adk/middlewares/modeltimeout/modeltimeout.go @@ -0,0 +1,54 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package modeltimeout provides ChatModelAgent middleware for enforcing model +// call and stream timeouts. +package modeltimeout + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// Middleware wraps model calls with timeout enforcement. +type Middleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + config *Config +} + +// New creates timeout middleware for the default *schema.Message ChatModelAgent. +func New(config *Config) adk.ChatModelAgentMiddleware { + return NewTyped[*schema.Message](config) +} + +// NewTyped creates timeout middleware for a typed ChatModelAgent. +func NewTyped[M adk.MessageType](config *Config) adk.TypedChatModelAgentMiddleware[M] { + return &Middleware[M]{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[M]{}, + config: config, + } +} + +// WrapModel installs timeout enforcement around the next model. +func (m *Middleware[M]) WrapModel(_ context.Context, next model.BaseModel[M], _ *adk.TypedModelContext[M]) (model.BaseModel[M], error) { + if !IsConfigActive(m.config) { + return next, nil + } + return NewTypedTimeoutModelWrapper(next, m.config), nil +} diff --git a/adk/middlewares/modeltimeout/modeltimeout_test.go b/adk/middlewares/modeltimeout/modeltimeout_test.go new file mode 100644 index 000000000..9ef771fb3 --- /dev/null +++ b/adk/middlewares/modeltimeout/modeltimeout_test.go @@ -0,0 +1,62 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package modeltimeout + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type blockingChatModel struct{} + +func (m *blockingChatModel) Generate(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + <-ctx.Done() + return nil, ctx.Err() +} + +func (m *blockingChatModel) Stream(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + <-ctx.Done() + return nil, ctx.Err() +} + +func TestMiddlewareWrapModel(t *testing.T) { + mw := New(&Config{CallTimeout: 10 * time.Millisecond}) + wrapped, err := mw.WrapModel(context.Background(), &blockingChatModel{}, nil) + require.NoError(t, err) + + _, err = wrapped.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.ErrorIs(t, err, ErrModelTimeout) + + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseCall, timeoutErr.Phase) +} + +func TestInactiveMiddlewareDelegates(t *testing.T) { + m := &blockingChatModel{} + + mw := New(&Config{}) + wrapped, err := mw.WrapModel(context.Background(), m, nil) + require.NoError(t, err) + require.Same(t, m, wrapped) +} diff --git a/adk/middlewares/modeltimeout/timeout.go b/adk/middlewares/modeltimeout/timeout.go new file mode 100644 index 000000000..696e44b31 --- /dev/null +++ b/adk/middlewares/modeltimeout/timeout.go @@ -0,0 +1,498 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package modeltimeout + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + "sync/atomic" + "time" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// Phase identifies which part of a model call exceeded its budget. +type Phase string + +const ( + // PhaseCall means Generate or Stream opening exceeded its budget. + PhaseCall Phase = "call" + // PhaseFirstChunk means no stream chunk arrived before the first-chunk budget. + PhaseFirstChunk Phase = "first_chunk" + // PhaseStreamIdle means the stream exceeded its inter-chunk idle budget. + PhaseStreamIdle Phase = "stream_idle" + // PhaseTotal means the whole Generate call or Stream lifecycle exceeded its budget. + PhaseTotal Phase = "total" +) + +// ErrModelTimeout is the sentinel matched by Error. +var ErrModelTimeout = errors.New("model timeout") + +// Config configures opt-in timeout enforcement for ChatModel calls. +// +// Timeout errors are surfaced as *Error and can be handled by +// ModelRetryConfig.ShouldRetry/IsRetryAble and ModelFailoverConfig.ShouldFailover. +// If nil or all durations are <= 0, no timeout wrapper is installed. +// +// Timeouts are per model attempt because the timeout wrapper sits inside retry +// and failover. Providers must respect context cancellation for Generate/Stream +// opening and context cancellation or StreamReader.Close for stream-body cleanup. +type Config struct { + // CallTimeout bounds Generate and Stream until Stream returns a reader. + // For Generate this is effectively the non-streaming model call timeout. + // For Stream this is the request-open/header/reader-acquisition timeout. + CallTimeout time.Duration + + // FirstChunkTimeout bounds the time from Stream returning a reader to the + // first successful chunk. + FirstChunkTimeout time.Duration + + // StreamIdleTimeout bounds the gap between successful stream chunks after + // the first chunk. + StreamIdleTimeout time.Duration + + // TotalTimeout bounds the whole Generate call or whole Stream lifecycle. + // It is per model attempt when retry/failover are configured. + TotalTimeout time.Duration +} + +// Error reports a model timeout without prescribing retry policy. +type Error struct { + Phase Phase + Timeout time.Duration + Elapsed time.Duration + ChunksReceived int +} + +func (e *Error) Error() string { + if e == nil { + return ErrModelTimeout.Error() + } + return fmt.Sprintf("model timeout: phase=%s timeout=%s elapsed=%s chunks_received=%d", + e.Phase, e.Timeout, e.Elapsed, e.ChunksReceived) +} + +func (e *Error) Is(target error) bool { + return target == ErrModelTimeout +} + +// IsModelTimeoutBeforeOutput reports whether the timeout happened before any +// stream output reached downstream consumers. +func (e *Error) IsModelTimeoutBeforeOutput() bool { + return e != nil && e.ChunksReceived == 0 +} + +// ModelTimeoutSpanMeta exposes timeout details to packages that should not +// import this middleware package directly. +func (e *Error) ModelTimeoutSpanMeta() (phase string, timeout time.Duration, elapsed time.Duration, chunksReceived int) { + if e == nil { + return "", 0, 0, 0 + } + return string(e.Phase), e.Timeout, e.Elapsed, e.ChunksReceived +} + +// AsModelTimeout extracts a Error from err. +func AsModelTimeout(err error) (*Error, bool) { + var timeoutErr *Error + if errors.As(err, &timeoutErr) { + return timeoutErr, true + } + return nil, false +} + +// IsModelTimeoutBeforeOutput reports whether err is a timeout that happened +// before any stream output reached downstream consumers. +func IsModelTimeoutBeforeOutput(err error) bool { + timeoutErr, ok := AsModelTimeout(err) + return ok && timeoutErr.ChunksReceived == 0 +} + +func init() { + schema.RegisterName[*Error]("_eino_adk_model_timeout_error") +} + +type typedTimeoutModelWrapper[M adk.MessageType] struct { + inner model.BaseModel[M] + config *Config +} + +// NewTypedTimeoutModelWrapper wraps a model with timeout enforcement. +// +// Prefer configuring this through adk/middlewares/modeltimeout so timeout +// behavior composes with other ChatModelAgent middlewares. +func NewTypedTimeoutModelWrapper[M adk.MessageType](inner model.BaseModel[M], config *Config) model.BaseModel[M] { + return &typedTimeoutModelWrapper[M]{inner: inner, config: config} +} + +func newTypedTimeoutModelWrapper[M adk.MessageType](inner model.BaseModel[M], config *Config) model.BaseModel[M] { + return NewTypedTimeoutModelWrapper(inner, config) +} + +// IsConfigActive reports whether config enables any timeout. +func IsConfigActive(config *Config) bool { + return config != nil && (config.CallTimeout > 0 || + config.FirstChunkTimeout > 0 || + config.StreamIdleTimeout > 0 || + config.TotalTimeout > 0) +} + +func isConfigActive(config *Config) bool { + return IsConfigActive(config) +} + +func minPositiveTimeout(callTimeout, totalTimeout time.Duration) (time.Duration, Phase, bool) { + switch { + case callTimeout > 0 && totalTimeout > 0: + if totalTimeout <= callTimeout { + return totalTimeout, PhaseTotal, true + } + return callTimeout, PhaseCall, true + case callTimeout > 0: + return callTimeout, PhaseCall, true + case totalTimeout > 0: + return totalTimeout, PhaseTotal, true + default: + return 0, "", false + } +} + +func modelTimeoutError(phase Phase, timeout time.Duration, started time.Time, chunks int) *Error { + return &Error{ + Phase: phase, + Timeout: timeout, + Elapsed: time.Since(started), + ChunksReceived: chunks, + } +} + +type timeoutGenerateResult[M adk.MessageType] struct { + msg M + err error +} + +func (w *typedTimeoutModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + timeout, phase, ok := minPositiveTimeout(w.config.CallTimeout, w.config.TotalTimeout) + if !ok { + return w.inner.Generate(ctx, input, opts...) + } + + started := time.Now() + timeoutCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + resultCh := make(chan timeoutGenerateResult[M], 1) + go func() { + msg, err := w.inner.Generate(timeoutCtx, input, opts...) + resultCh <- timeoutGenerateResult[M]{msg: msg, err: err} + }() + + select { + case result := <-resultCh: + if ctx.Err() == nil && errors.Is(result.err, context.DeadlineExceeded) && errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) { + var zero M + return zero, modelTimeoutError(phase, timeout, started, 0) + } + return result.msg, result.err + case <-ctx.Done(): + var zero M + cancel() + return zero, ctx.Err() + case <-timeoutCtx.Done(): + var zero M + cancel() + if ctx.Err() != nil { + return zero, ctx.Err() + } + return zero, modelTimeoutError(phase, timeout, started, 0) + } +} + +type timeoutStreamOpenResult[M adk.MessageType] struct { + reader *schema.StreamReader[M] + err error +} + +func (w *typedTimeoutModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + if !isConfigActive(w.config) { + return w.inner.Stream(ctx, input, opts...) + } + + started := time.Now() + bodyTimeoutActive := w.hasStreamBodyTimeout() + streamCtx := ctx + cancel := func() {} + if bodyTimeoutActive && w.config.TotalTimeout > 0 { + streamCtx, cancel = newStreamOpenTimeoutContext(ctx, w.config.TotalTimeout) + } else if bodyTimeoutActive || w.config.CallTimeout > 0 { + streamCtx, cancel = newStreamOpenCancelContext(ctx) + } + + resultCh := make(chan timeoutStreamOpenResult[M], 1) + done := make(chan struct{}) + accepted := make(chan struct{}) + go func() { + reader, err := w.inner.Stream(streamCtx, input, opts...) + result := timeoutStreamOpenResult[M]{reader: reader, err: err} + select { + case <-done: + if reader != nil { + reader.Close() + } + case resultCh <- result: + select { + case <-accepted: + case <-done: + if reader != nil { + reader.Close() + } + } + } + }() + + openTimeout, openPhase, hasOpenTimeout := minPositiveTimeout(w.config.CallTimeout, w.config.TotalTimeout) + var openTimer *time.Timer + var openTimeoutCh <-chan time.Time + if hasOpenTimeout { + openTimer = time.NewTimer(openTimeout) + openTimeoutCh = openTimer.C + defer openTimer.Stop() + } + + var result timeoutStreamOpenResult[M] + select { + case result = <-resultCh: + close(accepted) + if ctx.Err() == nil && hasOpenTimeout && result.err != nil && + streamCtx.Err() != nil && time.Since(started) >= openTimeout { + cancel() + return nil, modelTimeoutError(openPhase, openTimeout, started, 0) + } + if result.err != nil { + cancel() + return nil, result.err + } + case <-ctx.Done(): + close(done) + cancel() + return nil, ctx.Err() + case <-openTimeoutCh: + close(done) + cancel() + return nil, modelTimeoutError(openPhase, openTimeout, started, 0) + case <-streamCtx.Done(): + close(done) + cancel() + if ctx.Err() != nil { + return nil, ctx.Err() + } + return nil, modelTimeoutError(PhaseTotal, w.config.TotalTimeout, started, 0) + } + + if result.reader == nil { + cancel() + return nil, errors.New("model Stream returned nil reader without error") + } + if !bodyTimeoutActive { + return result.reader, nil + } + return w.wrapStreamBody(ctx, streamCtx, cancel, result.reader, started), nil +} + +func newStreamOpenCancelContext(ctx context.Context) (context.Context, context.CancelFunc) { + return context.WithCancel(ctx) +} + +func newStreamOpenTimeoutContext(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) { + return context.WithTimeout(ctx, timeout) +} + +func (w *typedTimeoutModelWrapper[M]) hasStreamBodyTimeout() bool { + return w.config.FirstChunkTimeout > 0 || w.config.StreamIdleTimeout > 0 || w.config.TotalTimeout > 0 +} + +type timeoutStreamWriter[M adk.MessageType] struct { + writer *schema.StreamWriter[M] + done chan struct{} + once sync.Once + mu sync.Mutex + closed bool +} + +func newTimeoutStreamWriter[M adk.MessageType](writer *schema.StreamWriter[M]) *timeoutStreamWriter[M] { + return &timeoutStreamWriter[M]{ + writer: writer, + done: make(chan struct{}), + } +} + +func (w *timeoutStreamWriter[M]) send(msg M, err error) bool { + w.mu.Lock() + defer w.mu.Unlock() + if w.closed { + return true + } + return w.writer.Send(msg, err) +} + +func (w *timeoutStreamWriter[M]) close() { + w.once.Do(func() { + w.mu.Lock() + w.closed = true + w.writer.Close() + w.mu.Unlock() + close(w.done) + }) +} + +func (w *typedTimeoutModelWrapper[M]) wrapStreamBody( + ctx context.Context, + streamCtx context.Context, + cancel context.CancelFunc, + upstream *schema.StreamReader[M], + started time.Time, +) *schema.StreamReader[M] { + reader, writer := schema.Pipe[M](1) + terminal := newTimeoutStreamWriter(writer) + var chunks int32 + activity := make(chan struct{}, 1) + var finishOnce sync.Once + + finish := func(err error) { + finishOnce.Do(func() { + if err != nil { + var zero M + terminal.send(zero, err) + } + terminal.close() + upstream.Close() + cancel() + }) + } + + go func() { + for { + msg, err := upstream.Recv() + if err == io.EOF { + finish(nil) + return + } + if err != nil { + if ctx.Err() != nil { + finish(ctx.Err()) + return + } + if streamCtx.Err() != nil && w.config.TotalTimeout > 0 { + finish(modelTimeoutError(PhaseTotal, w.config.TotalTimeout, started, int(atomic.LoadInt32(&chunks)))) + return + } + finish(err) + return + } + if terminal.send(msg, nil) { + finish(nil) + return + } + atomic.AddInt32(&chunks, 1) + select { + case activity <- struct{}{}: + default: + } + } + }() + + go func() { + firstReceived := false + var inactivityTimer *time.Timer + var inactivityCh <-chan time.Time + resetInactivity := func(d time.Duration) { + if inactivityTimer != nil { + if !inactivityTimer.Stop() { + select { + case <-inactivityTimer.C: + default: + } + } + } + if d > 0 { + inactivityTimer = time.NewTimer(d) + inactivityCh = inactivityTimer.C + } else { + inactivityCh = nil + } + } + defer func() { + if inactivityTimer != nil { + inactivityTimer.Stop() + } + }() + + resetInactivity(w.config.FirstChunkTimeout) + var totalTimer *time.Timer + var totalCh <-chan time.Time + if w.config.TotalTimeout > 0 { + remaining := time.Until(started.Add(w.config.TotalTimeout)) + if remaining < 0 { + remaining = 0 + } + totalTimer = time.NewTimer(remaining) + totalCh = totalTimer.C + defer totalTimer.Stop() + } + + for { + select { + case <-terminal.done: + return + case <-activity: + if !firstReceived { + firstReceived = true + } + resetInactivity(w.config.StreamIdleTimeout) + case <-inactivityCh: + phase := PhaseFirstChunk + timeout := w.config.FirstChunkTimeout + if firstReceived { + phase = PhaseStreamIdle + timeout = w.config.StreamIdleTimeout + } + finish(modelTimeoutError(phase, timeout, started, int(atomic.LoadInt32(&chunks)))) + return + case <-totalCh: + finish(modelTimeoutError(PhaseTotal, w.config.TotalTimeout, started, int(atomic.LoadInt32(&chunks)))) + return + case <-streamCtx.Done(): + if ctx.Err() != nil { + finish(ctx.Err()) + return + } + if w.config.TotalTimeout > 0 { + finish(modelTimeoutError(PhaseTotal, w.config.TotalTimeout, started, int(atomic.LoadInt32(&chunks)))) + return + } + finish(streamCtx.Err()) + return + } + } + }() + + return reader +} diff --git a/adk/middlewares/modeltimeout/timeout_test.go b/adk/middlewares/modeltimeout/timeout_test.go new file mode 100644 index 000000000..4ec730e44 --- /dev/null +++ b/adk/middlewares/modeltimeout/timeout_test.go @@ -0,0 +1,622 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package modeltimeout + +import ( + "context" + "errors" + "io" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + . "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type fakeChatModel struct { + callbacksEnabled bool + generate func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) + stream func(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) +} + +func (m *fakeChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return m.generate(ctx, input, opts...) +} + +func (m *fakeChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return m.stream(ctx, input, opts...) +} + +func (m *fakeChatModel) BindTools([]*schema.ToolInfo) error { + return nil +} + +func (m *fakeChatModel) IsCallbacksEnabled() bool { + return m.callbacksEnabled +} + +type mockAgenticModel struct { + generateFn func(context.Context, []*schema.AgenticMessage, ...model.Option) (*schema.AgenticMessage, error) + streamFn func(context.Context, []*schema.AgenticMessage, ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) +} + +func (m *mockAgenticModel) Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return m.generateFn(ctx, input, opts...) +} + +func (m *mockAgenticModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + if m.streamFn != nil { + return m.streamFn(ctx, input, opts...) + } + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.AgenticMessage{msg}), nil +} + +func instantBackoff(context.Context, int) time.Duration { + return 0 +} + +func drainTimeoutAgentEvents(iter *AsyncIterator[*AgentEvent]) []*AgentEvent { + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + return events + } + events = append(events, event) + } +} + +func contextAwareMessageStream(ctx context.Context, chunks ...*schema.Message) *schema.StreamReader[*schema.Message] { + reader, writer := schema.Pipe[*schema.Message](len(chunks) + 1) + for _, chunk := range chunks { + writer.Send(chunk, nil) + } + go func() { + <-ctx.Done() + writer.Send(nil, ctx.Err()) + }() + return reader +} + +func TestModelTimeoutGenerateCallTimeout(t *testing.T) { + release := make(chan struct{}) + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + <-release + return schema.AssistantMessage("late", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + defer close(release) + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{CallTimeout: 10 * time.Millisecond}) + started := time.Now() + _, err := wrapped.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.Less(t, time.Since(started), 200*time.Millisecond) + + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseCall, timeoutErr.Phase) + require.Equal(t, 0, timeoutErr.ChunksReceived) + require.True(t, errors.Is(err, ErrModelTimeout)) + require.True(t, IsModelTimeoutBeforeOutput(err)) +} + +func TestModelTimeoutGenerateParentCancellation(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{CallTimeout: time.Second}) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.ErrorIs(t, err, context.Canceled) + require.False(t, errors.Is(err, ErrModelTimeout)) +} + +func TestModelTimeoutStreamFirstChunkTimeout(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return contextAwareMessageStream(ctx), nil + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{FirstChunkTimeout: 10 * time.Millisecond}) + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + defer stream.Close() + + _, err = stream.Recv() + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseFirstChunk, timeoutErr.Phase) + require.Equal(t, 0, timeoutErr.ChunksReceived) + require.True(t, IsModelTimeoutBeforeOutput(err)) +} + +func TestModelTimeoutStreamOpenCooperativeTimeout(t *testing.T) { + cooperated := make(chan struct{}) + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + <-ctx.Done() + close(cooperated) + return nil, ctx.Err() + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{CallTimeout: 10 * time.Millisecond}) + _, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseCall, timeoutErr.Phase) + select { + case <-cooperated: + case <-time.After(time.Second): + t.Fatal("stream-open context was not canceled on timeout") + } +} + +func TestAttack_StreamOpenTimeoutDoesNotRequireProviderCooperation(t *testing.T) { + release := make(chan struct{}) + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + <-release + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("late", nil)}), nil + }, + } + defer close(release) + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{CallTimeout: 10 * time.Millisecond}) + errCh := make(chan error, 1) + go func() { + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + if stream != nil { + stream.Close() + } + errCh <- err + }() + + select { + case err := <-errCh: + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseCall, timeoutErr.Phase) + case <-time.After(200 * time.Millisecond): + t.Fatal("stream open did not return at CallTimeout when provider ignored context") + } +} + +func TestModelTimeoutStreamIdleTimeoutAfterOutput(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return contextAwareMessageStream(ctx, schema.AssistantMessage("first", nil)), nil + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{StreamIdleTimeout: 10 * time.Millisecond}) + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + defer stream.Close() + + msg, err := stream.Recv() + require.NoError(t, err) + require.Equal(t, "first", msg.Content) + + _, err = stream.Recv() + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseStreamIdle, timeoutErr.Phase) + require.Equal(t, 1, timeoutErr.ChunksReceived) + require.False(t, IsModelTimeoutBeforeOutput(err)) +} + +func TestModelTimeoutStreamTotalTimeoutAfterOutput(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return contextAwareMessageStream(ctx, schema.AssistantMessage("first", nil)), nil + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{TotalTimeout: 20 * time.Millisecond}) + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + defer stream.Close() + + msg, err := stream.Recv() + require.NoError(t, err) + require.Equal(t, "first", msg.Content) + + _, err = stream.Recv() + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseTotal, timeoutErr.Phase) + require.Equal(t, 1, timeoutErr.ChunksReceived) +} + +func TestAttack_StreamBodyTimeoutDoesNotRequireUpstreamRecvCooperation(t *testing.T) { + upstreamReader, upstreamWriter := schema.Pipe[*schema.Message](0) + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return upstreamReader, nil + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{FirstChunkTimeout: 10 * time.Millisecond}) + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + defer stream.Close() + defer upstreamWriter.Close() + + errCh := make(chan error, 1) + go func() { + _, recvErr := stream.Recv() + errCh <- recvErr + }() + + select { + case err := <-errCh: + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseFirstChunk, timeoutErr.Phase) + require.Equal(t, 0, timeoutErr.ChunksReceived) + case <-time.After(200 * time.Millisecond): + t.Fatal("stream body did not return at FirstChunkTimeout when upstream Recv stayed blocked") + } +} + +func TestModelTimeoutGenerateTotalBeatsCallTimeout(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{ + CallTimeout: time.Second, + TotalTimeout: 10 * time.Millisecond, + }) + _, err := wrapped.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseTotal, timeoutErr.Phase) +} + +func TestModelTimeoutStreamDownstreamCloseClosesUpstream(t *testing.T) { + upstreamReader, upstreamWriter := schema.Pipe[*schema.Message](0) + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return upstreamReader, nil + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{StreamIdleTimeout: time.Second}) + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + stream.Close() + + secondSent := make(chan bool, 1) + go func() { + secondSent <- upstreamWriter.Send(schema.AssistantMessage("second", nil), nil) + }() + select { + case <-secondSent: + case <-time.After(time.Second): + t.Fatal("wrapper did not receive the post-close upstream chunk") + } + + closed := make(chan bool, 1) + go func() { + closed <- upstreamWriter.Send(schema.AssistantMessage("third", nil), nil) + }() + select { + case got := <-closed: + require.True(t, got, "upstream reader should be closed after downstream close is observed") + case <-time.After(time.Second): + t.Fatal("upstream reader was not closed after downstream close") + } +} + +func TestModelTimeoutChatModelAgentRetryIntegration(t *testing.T) { + var calls int32 + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + if atomic.AddInt32(&calls, 1) == 1 { + <-ctx.Done() + return nil, ctx.Err() + } + return schema.AssistantMessage("success", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + agent, err := NewChatModelAgent(context.Background(), &ChatModelAgentConfig{ + Name: "timeout-retry", + Description: "timeout retry", + Model: m, + Handlers: []ChatModelAgentMiddleware{New(&Config{CallTimeout: 10 * time.Millisecond})}, + ModelRetryConfig: &ModelRetryConfig{MaxRetries: 1, BackoffFunc: instantBackoff}, + }) + require.NoError(t, err) + + events := drainTimeoutAgentEvents(agent.Run(context.Background(), &AgentInput{Messages: []Message{schema.UserMessage("hi")}})) + require.Len(t, events, 1) + require.NoError(t, events[0].Err) + require.Equal(t, "success", events[0].Output.MessageOutput.Message.Content) + require.Equal(t, int32(2), atomic.LoadInt32(&calls)) +} + +func TestModelTimeoutAgenticMessageGenerate(t *testing.T) { + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + } + wrapped := newTypedTimeoutModelWrapper[*schema.AgenticMessage](m, &Config{CallTimeout: 10 * time.Millisecond}) + + _, err := wrapped.Generate(context.Background(), []*schema.AgenticMessage{schema.UserAgenticMessage("hi")}) + timeoutErr, ok := AsModelTimeout(err) + require.True(t, ok) + require.Equal(t, PhaseCall, timeoutErr.Phase) +} + +func TestModelTimeoutTimelineEventContainsTimeoutMeta(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + agent, err := NewChatModelAgent(context.Background(), &ChatModelAgentConfig{ + Name: "timeout-timeline", + Description: "timeout timeline", + Model: m, + Handlers: []ChatModelAgentMiddleware{New(&Config{CallTimeout: 10 * time.Millisecond})}, + }) + require.NoError(t, err) + + var endEvent *SessionEvent[*schema.Message] + iter := agent.Run(context.Background(), &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == SessionEventSpanModelRequestEnd { + endEvent = event.SessionEventVariant.Event + } + } + require.NotNil(t, endEvent) + require.Equal(t, "error", endEvent.Span.Status) + require.Contains(t, endEvent.Span.Err, "model timeout") + require.NotNil(t, endEvent.Span.Model.Timeout) + require.Equal(t, string(PhaseCall), endEvent.Span.Model.Timeout.Phase) + +} + +func TestAttack_ModelTimeoutRetryExhaustionKeepsTimelineTimeoutMeta(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + agent, err := NewChatModelAgent(context.Background(), &ChatModelAgentConfig{ + Name: "timeout-retry-exhausted-timeline", + Description: "timeout retry exhausted timeline", + Model: m, + Handlers: []ChatModelAgentMiddleware{New(&Config{CallTimeout: 10 * time.Millisecond})}, + ModelRetryConfig: &ModelRetryConfig{MaxRetries: 0, BackoffFunc: instantBackoff}, + }) + require.NoError(t, err) + + var endEvent *SessionEvent[*schema.Message] + iter := agent.Run(context.Background(), &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == SessionEventSpanModelRequestEnd { + endEvent = event.SessionEventVariant.Event + } + } + require.NotNil(t, endEvent) + require.Contains(t, endEvent.Span.Err, "model timeout") + require.NotNil(t, endEvent.Span.Model.Timeout) + require.Equal(t, string(PhaseCall), endEvent.Span.Model.Timeout.Phase) +} + +func TestModelTimeoutHelperContracts(t *testing.T) { + var nilTimeout *Error + require.Equal(t, ErrModelTimeout.Error(), nilTimeout.Error()) + + timeoutErr := &Error{ + Phase: PhaseStreamIdle, + Timeout: time.Second, + Elapsed: time.Millisecond, + ChunksReceived: 2, + } + require.ErrorIs(t, timeoutErr, ErrModelTimeout) + require.Contains(t, timeoutErr.Error(), "chunks_received=2") + + extracted, ok := AsModelTimeout(timeoutErr) + require.True(t, ok) + require.Same(t, timeoutErr, extracted) + require.False(t, IsModelTimeoutBeforeOutput(timeoutErr)) + require.True(t, IsModelTimeoutBeforeOutput(&Error{ChunksReceived: 0})) + + _, ok = AsModelTimeout(io.EOF) + require.False(t, ok) + require.False(t, isConfigActive(nil)) + require.False(t, isConfigActive(&Config{})) + require.True(t, isConfigActive(&Config{StreamIdleTimeout: time.Second})) + + timeout, phase, ok := minPositiveTimeout(0, 0) + require.False(t, ok) + require.Zero(t, timeout) + require.Empty(t, phase) +} + +func TestModelTimeoutInactiveConfigDelegates(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("generated", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("streamed", nil)}), nil + }, + } + + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{}) + msg, err := wrapped.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "generated", msg.Content) + + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + defer stream.Close() + + chunk, err := stream.Recv() + require.NoError(t, err) + require.Equal(t, "streamed", chunk.Content) + _, err = stream.Recv() + require.ErrorIs(t, err, io.EOF) +} + +func TestModelTimeoutStreamOpenErrorPaths(t *testing.T) { + t.Run("nil reader without error", func(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, nil + }, + } + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{FirstChunkTimeout: time.Second}) + + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, stream) + require.Error(t, err) + require.Contains(t, err.Error(), "nil reader") + }) + + t.Run("provider error passes through", func(t *testing.T) { + providerErr := errors.New("provider stream failed") + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, providerErr + }, + } + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{CallTimeout: time.Second}) + + stream, err := wrapped.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, stream) + require.ErrorIs(t, err, providerErr) + }) + + t.Run("parent cancellation wins open", func(t *testing.T) { + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + <-ctx.Done() + return nil, ctx.Err() + }, + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + wrapped := newTypedTimeoutModelWrapper[*schema.Message](m, &Config{CallTimeout: time.Second}) + + stream, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, stream) + require.ErrorIs(t, err, context.Canceled) + require.False(t, errors.Is(err, ErrModelTimeout)) + }) +} diff --git a/adk/middlewares/patchtoolcalls/patchtoolcalls.go b/adk/middlewares/patchtoolcalls/patchtoolcalls.go index 484c8811f..06e85d0c2 100644 --- a/adk/middlewares/patchtoolcalls/patchtoolcalls.go +++ b/adk/middlewares/patchtoolcalls/patchtoolcalls.go @@ -20,12 +20,15 @@ package patchtoolcalls import ( "context" "fmt" + "strings" "github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk/internal" "github.com/cloudwego/eino/schema" ) +const syntheticAgenticToolResultMarker = "_eino_patch_tool_calls_synthetic" + // Config defines the configuration options for the patch tool calls middleware. type Config struct { // PatchedContentGenerator is an optional custom function to generate the content @@ -40,6 +43,22 @@ type Config struct { // - string: the content to use for the patched tool message // - error: any error that occurred during generation PatchedContentGenerator func(ctx context.Context, toolName, toolCallID string) (string, error) + + // RemoveOrphanResults removes tool result messages or result blocks whose call ID + // does not match any previous assistant tool call. Disabled by default. + RemoveOrphanResults bool + + // RemoveDuplicateResults removes duplicate tool result messages or result blocks + // after the first result kept for a call ID. Disabled by default. + RemoveDuplicateResults bool + + // Strict validates the history and returns an error without mutating state when + // missing, orphan, duplicate, or empty-ID mismatches are found. Disabled by default. + Strict bool + + // MarkSynthetic marks generated AgenticMessage tool results in Extra so callers + // can identify mechanical repairs. Disabled by default. + MarkSynthetic bool } // NewTyped creates a new generic patch tool calls middleware. @@ -50,8 +69,9 @@ func NewTyped[M adk.MessageType](_ context.Context, cfg *Config) (adk.TypedChatM if cfg == nil { cfg = &Config{} } + cfgCopy := *cfg return &typedMiddleware[M]{ - gen: cfg.PatchedContentGenerator, + cfg: cfgCopy, }, nil } @@ -65,7 +85,7 @@ func New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) type typedMiddleware[M adk.MessageType] struct { *adk.TypedBaseChatModelAgentMiddleware[M] - gen func(ctx context.Context, toolName, toolCallID string) (string, error) + cfg Config } func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], @@ -78,93 +98,388 @@ func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state var zero M switch any(zero).(type) { case *schema.Message: - return patchToolCallsForMessage(ctx, m.gen, any(state).(*adk.TypedChatModelAgentState[*schema.Message]), mc) + return patchToolCallsForMessage(ctx, m.cfg, any(state).(*adk.TypedChatModelAgentState[*schema.Message]), mc) case *schema.AgenticMessage: - return patchToolCallsForAgenticMessage(ctx, m.gen, any(state).(*adk.TypedChatModelAgentState[*schema.AgenticMessage]), mc) + return patchToolCallsForAgenticMessage(ctx, m.cfg, any(state).(*adk.TypedChatModelAgentState[*schema.AgenticMessage]), mc) default: panic("unreachable: unknown MessageType") } } func patchToolCallsForMessage[M adk.MessageType](ctx context.Context, - gen func(ctx context.Context, toolName, toolCallID string) (string, error), + cfg Config, state *adk.TypedChatModelAgentState[*schema.Message], - _ *adk.TypedModelContext[M], -) (context.Context, *adk.TypedChatModelAgentState[M], error) { - patched := make([]*schema.Message, 0, len(state.Messages)) + _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) { + + plan, err := buildMessageNormalizationPlan(ctx, cfg, state.Messages) + if err != nil { + return ctx, nil, err + } + if err := sendNormalizationEvents(ctx, plan.events); err != nil { + return ctx, nil, err + } + + nState := *state + nState.Messages = plan.messages + return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil +} + +func patchToolCallsForAgenticMessage[M adk.MessageType](ctx context.Context, + cfg Config, + state *adk.TypedChatModelAgentState[*schema.AgenticMessage], + _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) { + + plan, err := buildAgenticNormalizationPlan(ctx, cfg, state.Messages) + if err != nil { + return ctx, nil, err + } + if err := sendNormalizationEvents(ctx, plan.events); err != nil { + return ctx, nil, err + } + + nState := *state + nState.Messages = plan.messages + return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil +} + +type mismatchCounts struct { + missing int + orphan int + duplicate int + emptyID int +} + +func (c mismatchCounts) hasMismatch() bool { + return c.missing > 0 || c.orphan > 0 || c.duplicate > 0 || c.emptyID > 0 +} + +func (c mismatchCounts) strictError() error { + return fmt.Errorf("patchtoolcalls strict validation failed: missing=%d orphan=%d duplicate=%d empty_tool_call_id=%d", + c.missing, c.orphan, c.duplicate, c.emptyID) +} + +type normalizationPlan[M adk.MessageType] struct { + messages []M + events []*adk.SessionEvent[M] + counts mismatchCounts +} + +func buildMessageNormalizationPlan(ctx context.Context, cfg Config, messages []*schema.Message) (*normalizationPlan[*schema.Message], error) { + ensureMessageIDs(messages) + + counts := analyzeMessages(messages) + if cfg.Strict && counts.hasMismatch() { + return nil, counts.strictError() + } - for i, msg := range state.Messages { - patched = append(patched, msg) + keep := keptMessages(messages, cfg) + patched := make([]*schema.Message, 0, len(messages)+counts.missing) + inserted := make([]*adk.SessionEvent[*schema.Message], 0, counts.missing) + for i, msg := range messages { + if keep[i] { + patched = append(patched, msg) + } if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 { continue } - for _, tc := range msg.ToolCalls { - if hasCorrespondingToolMessage(state.Messages[i+1:], tc.ID) { + if tc.ID == "" || hasCorrespondingToolMessage(messages[i+1:], tc.ID) { continue } - - toolMsg, err := createPatchedToolMessage(ctx, gen, tc) + toolMsg, err := createPatchedToolMessage(ctx, cfg.PatchedContentGenerator, tc) if err != nil { - return ctx, nil, err + return nil, err } + adk.EnsureMessageID(toolMsg) patched = append(patched, toolMsg) + inserted = append(inserted, &adk.SessionEvent[*schema.Message]{ + Kind: adk.SessionEventMessageInserted, + MessageInserted: &adk.MessageInsertedEvent[*schema.Message]{ + Message: toolMsg, + BeforeMessageID: firstKeptMessageID(messages, keep, i+1), + }, + }) } } - nState := *state - nState.Messages = patched - return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil + events := make([]*adk.SessionEvent[*schema.Message], 0, len(inserted)+1) + events = append(events, inserted...) + if deletedIDs := deletedMessageIDs(messages, keep); len(deletedIDs) > 0 { + events = append(events, &adk.SessionEvent[*schema.Message]{ + Kind: adk.SessionEventMessagesDeleted, + MessagesDeleted: &adk.MessagesDeletedEvent{ + MessageIDs: deletedIDs, + }, + }) + } + + return &normalizationPlan[*schema.Message]{messages: patched, events: events, counts: counts}, nil } -func patchToolCallsForAgenticMessage[M adk.MessageType](ctx context.Context, - gen func(ctx context.Context, toolName, toolCallID string) (string, error), - state *adk.TypedChatModelAgentState[*schema.AgenticMessage], - _ *adk.TypedModelContext[M], -) (context.Context, *adk.TypedChatModelAgentState[M], error) { - patched := make([]*schema.AgenticMessage, 0, len(state.Messages)) +func analyzeMessages(messages []*schema.Message) mismatchCounts { + var counts mismatchCounts + previousCalls := make(map[string]struct{}) + seenResults := make(map[string]struct{}) + + for i, msg := range messages { + if msg.Role == schema.Tool { + if _, ok := previousCalls[msg.ToolCallID]; !ok { + counts.orphan++ + } else if _, ok := seenResults[msg.ToolCallID]; ok { + counts.duplicate++ + } else { + seenResults[msg.ToolCallID] = struct{}{} + } + } + if msg.Role != schema.Assistant { + continue + } + for _, tc := range msg.ToolCalls { + if tc.ID == "" { + counts.emptyID++ + continue + } + previousCalls[tc.ID] = struct{}{} + if !hasCorrespondingToolMessage(messages[i+1:], tc.ID) { + counts.missing++ + } + } + } - for i, msg := range state.Messages { - patched = append(patched, msg) + return counts +} - if msg.Role != schema.AgenticRoleTypeAssistant { +func ensureMessageIDs[M adk.MessageType](messages []M) { + for _, msg := range messages { + adk.EnsureMessageID(msg) + } +} + +func keptMessages(messages []*schema.Message, cfg Config) []bool { + keep := make([]bool, len(messages)) + previousCalls := make(map[string]struct{}) + seenResults := make(map[string]struct{}) + + for i, msg := range messages { + keep[i] = true + if msg.Role == schema.Tool { + _, valid := previousCalls[msg.ToolCallID] + _, duplicate := seenResults[msg.ToolCallID] + if !valid && cfg.RemoveOrphanResults { + keep[i] = false + } else if valid && duplicate && cfg.RemoveDuplicateResults { + keep[i] = false + } + if valid && !duplicate { + seenResults[msg.ToolCallID] = struct{}{} + } + } + if msg.Role != schema.Assistant { continue } + for _, tc := range msg.ToolCalls { + if tc.ID != "" { + previousCalls[tc.ID] = struct{}{} + } + } + } + + return keep +} + +func buildAgenticNormalizationPlan(ctx context.Context, cfg Config, messages []*schema.AgenticMessage) (*normalizationPlan[*schema.AgenticMessage], error) { + ensureMessageIDs(messages) + + counts := analyzeAgenticMessages(messages) + if cfg.Strict && counts.hasMismatch() { + return nil, counts.strictError() + } + + rewrites := agenticMessageRewrites(messages, cfg) + patched := make([]*schema.AgenticMessage, 0, len(messages)+counts.missing) + inserted := make([]*adk.SessionEvent[*schema.AgenticMessage], 0, counts.missing) + updated := make([]*adk.SessionEvent[*schema.AgenticMessage], 0) - // Collect tool call IDs from this assistant message. - var toolCalls []struct { - callID string - name string + for i, msg := range messages { + rewrite := rewrites[i] + if rewrite.keep { + patched = append(patched, rewrite.message) + if rewrite.updated { + updated = append(updated, &adk.SessionEvent[*schema.AgenticMessage]{ + Kind: adk.SessionEventMessageUpdated, + MessageUpdated: &adk.MessageUpdatedEvent[*schema.AgenticMessage]{ + MessageID: adk.GetMessageID(msg), + Message: rewrite.message, + }, + }) + } + } + if msg.Role != schema.AgenticRoleTypeAssistant { + continue } + for _, tc := range collectAgenticToolCalls(msg) { + if tc.callID == "" || hasCorrespondingAgenticToolResult(messages[i+1:], tc.callID) { + continue + } + toolMsg, err := createPatchedAgenticToolMessage(ctx, cfg.PatchedContentGenerator, tc.name, tc.callID) + if err != nil { + return nil, err + } + if cfg.MarkSynthetic { + markSyntheticAgenticToolResult(toolMsg) + } + adk.EnsureMessageID(toolMsg) + patched = append(patched, toolMsg) + inserted = append(inserted, &adk.SessionEvent[*schema.AgenticMessage]{ + Kind: adk.SessionEventMessageInserted, + MessageInserted: &adk.MessageInsertedEvent[*schema.AgenticMessage]{ + Message: toolMsg, + BeforeMessageID: firstKeptAgenticMessageID(messages, rewrites, i+1), + }, + }) + } + } + + events := make([]*adk.SessionEvent[*schema.AgenticMessage], 0, len(inserted)+len(updated)+1) + events = append(events, inserted...) + events = append(events, updated...) + if deletedIDs := deletedAgenticMessageIDs(messages, rewrites); len(deletedIDs) > 0 { + events = append(events, &adk.SessionEvent[*schema.AgenticMessage]{ + Kind: adk.SessionEventMessagesDeleted, + MessagesDeleted: &adk.MessagesDeletedEvent{ + MessageIDs: deletedIDs, + }, + }) + } + + return &normalizationPlan[*schema.AgenticMessage]{messages: patched, events: events, counts: counts}, nil +} + +type agenticToolCall struct { + callID string + name string +} + +type agenticRewrite struct { + message *schema.AgenticMessage + keep bool + updated bool +} + +func analyzeAgenticMessages(messages []*schema.AgenticMessage) mismatchCounts { + var counts mismatchCounts + previousCalls := make(map[string]struct{}) + seenResults := make(map[string]struct{}) + + for i, msg := range messages { for _, block := range msg.ContentBlocks { - if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil { - toolCalls = append(toolCalls, struct { - callID string - name string - }{callID: block.FunctionToolCall.CallID, name: block.FunctionToolCall.Name}) + callID, ok := agenticResultCallID(block) + if !ok { + continue + } + if _, valid := previousCalls[callID]; !valid { + counts.orphan++ + } else if _, duplicate := seenResults[callID]; duplicate { + counts.duplicate++ + } else { + seenResults[callID] = struct{}{} } } - if len(toolCalls) == 0 { + if msg.Role != schema.AgenticRoleTypeAssistant { continue } + for _, tc := range collectAgenticToolCalls(msg) { + if tc.callID == "" { + counts.emptyID++ + continue + } + previousCalls[tc.callID] = struct{}{} + if !hasCorrespondingAgenticToolResult(messages[i+1:], tc.callID) { + counts.missing++ + } + } + } + + return counts +} + +func agenticMessageRewrites(messages []*schema.AgenticMessage, cfg Config) []agenticRewrite { + rewrites := make([]agenticRewrite, len(messages)) + previousCalls := make(map[string]struct{}) + seenResults := make(map[string]struct{}) + + for i, msg := range messages { + rewrite := agenticRewrite{message: msg, keep: true} + blocks := make([]*schema.ContentBlock, 0, len(msg.ContentBlocks)) + removedBlock := false - for _, tc := range toolCalls { - if hasCorrespondingAgenticToolResult(state.Messages[i+1:], tc.callID) { + for _, block := range msg.ContentBlocks { + callID, ok := agenticResultCallID(block) + if !ok { + blocks = append(blocks, block) continue } + _, valid := previousCalls[callID] + _, duplicate := seenResults[callID] + remove := (!valid && cfg.RemoveOrphanResults) || (valid && duplicate && cfg.RemoveDuplicateResults) + if remove { + removedBlock = true + } else { + blocks = append(blocks, block) + } + if valid && !duplicate { + seenResults[callID] = struct{}{} + } + } - toolMsg, err := createPatchedAgenticToolMessage(ctx, gen, tc.name, tc.callID) - if err != nil { - return ctx, nil, err + if removedBlock { + if len(blocks) == 0 { + rewrite.keep = false + } else { + adk.EnsureMessageID(msg) + cp := *msg + cp.ContentBlocks = blocks + cp.Extra = copyStringAnyMap(msg.Extra) + rewrite.message = &cp + rewrite.updated = true + } + } + + if msg.Role == schema.AgenticRoleTypeAssistant { + for _, tc := range collectAgenticToolCalls(msg) { + if tc.callID != "" { + previousCalls[tc.callID] = struct{}{} + } } - patched = append(patched, toolMsg) } + rewrites[i] = rewrite } - nState := *state - nState.Messages = patched - return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil + return rewrites +} + +func collectAgenticToolCalls(msg *schema.AgenticMessage) []agenticToolCall { + toolCalls := make([]agenticToolCall, 0) + for _, block := range msg.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil { + toolCalls = append(toolCalls, agenticToolCall{callID: block.FunctionToolCall.CallID, name: block.FunctionToolCall.Name}) + } + } + return toolCalls +} + +func agenticResultCallID(block *schema.ContentBlock) (string, bool) { + if block == nil { + return "", false + } + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + return block.FunctionToolResult.CallID, true + } + if block.Type == schema.ContentBlockTypeToolSearchResult && block.ToolSearchFunctionToolResult != nil { + return block.ToolSearchFunctionToolResult.CallID, true + } + return "", false } func hasCorrespondingToolMessage(messages []*schema.Message, toolCallID string) bool { @@ -188,18 +503,10 @@ func hasCorrespondingAgenticToolResult(messages []*schema.AgenticMessage, toolCa } hasToolResult := false for _, block := range msg.ContentBlocks { - if block == nil { - continue - } - if block.Type == schema.ContentBlockTypeFunctionToolResult { + callID, ok := agenticResultCallID(block) + if ok { hasToolResult = true - if block.FunctionToolResult != nil && block.FunctionToolResult.CallID == toolCallID { - return true - } - } - if block.Type == schema.ContentBlockTypeToolSearchResult { - hasToolResult = true - if block.ToolSearchFunctionToolResult != nil && block.ToolSearchFunctionToolResult.CallID == toolCallID { + if callID == toolCallID { return true } } @@ -211,6 +518,87 @@ func hasCorrespondingAgenticToolResult(messages []*schema.AgenticMessage, toolCa return false } +func firstKeptMessageID(messages []*schema.Message, keep []bool, start int) string { + for i := start; i < len(messages); i++ { + if keep[i] { + adk.EnsureMessageID(messages[i]) + return adk.GetMessageID(messages[i]) + } + } + return "" +} + +func firstKeptAgenticMessageID(messages []*schema.AgenticMessage, rewrites []agenticRewrite, start int) string { + for i := start; i < len(messages); i++ { + if rewrites[i].keep { + adk.EnsureMessageID(messages[i]) + return adk.GetMessageID(messages[i]) + } + } + return "" +} + +func deletedMessageIDs(messages []*schema.Message, keep []bool) []string { + ids := make([]string, 0) + for i, msg := range messages { + if keep[i] { + continue + } + adk.EnsureMessageID(msg) + ids = append(ids, adk.GetMessageID(msg)) + } + return ids +} + +func deletedAgenticMessageIDs(messages []*schema.AgenticMessage, rewrites []agenticRewrite) []string { + ids := make([]string, 0) + for i, msg := range messages { + if rewrites[i].keep { + continue + } + adk.EnsureMessageID(msg) + ids = append(ids, adk.GetMessageID(msg)) + } + return ids +} + +func sendNormalizationEvents[M adk.MessageType](ctx context.Context, events []*adk.SessionEvent[M]) error { + for _, event := range events { + err := adk.TypedSendEvent(ctx, &adk.TypedAgentEvent[M]{ + SessionEventVariant: &adk.SessionEventVariant[M]{Event: event}, + }) + if isOutOfRunContextError(err) { + continue + } + if err != nil { + return err + } + } + return nil +} + +func isOutOfRunContextError(err error) bool { + return err != nil && strings.Contains(err.Error(), "must be called within a ChatModelAgent Run() or Resume() execution context") +} + +func markSyntheticAgenticToolResult(msg *schema.AgenticMessage) { + if msg.Extra == nil { + msg.Extra = make(map[string]any, 1) + } + msg.Extra[syntheticAgenticToolResultMarker] = true +} + +func copyStringAnyMap(src map[string]any) map[string]any { + if src == nil { + return nil + } + dst := make(map[string]any, len(src)) + for k, v := range src { + dst[k] = v + } + return dst +} + func createPatchedToolMessage(ctx context.Context, gen func(ctx context.Context, toolName, toolCallID string) (string, error), tc schema.ToolCall) (*schema.Message, error) { if gen != nil { content, err := gen(ctx, tc.Function.Name, tc.ID) diff --git a/adk/middlewares/patchtoolcalls/patchtoolcalls_test.go b/adk/middlewares/patchtoolcalls/patchtoolcalls_test.go index 2fdb3c1c3..c098ef37a 100644 --- a/adk/middlewares/patchtoolcalls/patchtoolcalls_test.go +++ b/adk/middlewares/patchtoolcalls/patchtoolcalls_test.go @@ -153,6 +153,31 @@ func assertToolResultName[M adk.MessageType](t *testing.T, msg M, expectedName s } } +func collectToolResultIDs[M adk.MessageType](messages []M) []string { + var ids []string + for _, msg := range messages { + switch m := any(msg).(type) { + case *schema.Message: + if m.Role == schema.Tool { + ids = append(ids, m.ToolCallID) + } + case *schema.AgenticMessage: + for _, block := range m.ContentBlocks { + if callID, ok := agenticResultCallID(block); ok { + ids = append(ids, callID) + } + } + } + } + return ids +} + +func assertSyntheticMarker(t *testing.T, msg *schema.AgenticMessage, expected bool) { + t.Helper() + v, ok := msg.Extra[syntheticAgenticToolResultMarker] + assert.Equal(t, expected, ok && v == true) +} + func testPatchToolCallsGeneric[M adk.MessageType](t *testing.T) { ctx := context.Background() @@ -282,6 +307,182 @@ func TestPatchToolCallsGeneric(t *testing.T) { t.Run("AgenticMessage", testPatchToolCallsGeneric[*schema.AgenticMessage]) } +func testPatchToolCallsRemoveOrphanResults[M adk.MessageType](t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{RemoveOrphanResults: true}) + require.NoError(t, err) + + state := &adk.TypedChatModelAgentState[M]{Messages: []M{ + makeToolResultMsg[M]("orphan", "call_orphan", "tool_orphan"), + makeAssistantMsgWithToolCalls[M]("", []testToolCall{{ID: "call_1", Name: "tool_a", Arguments: "{}"}}), + makeToolResultMsg[M]("result", "call_1", "tool_a"), + }} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + assert.Equal(t, []string{"call_1"}, collectToolResultIDs(newState.Messages)) +} + +func TestPatchToolCallsRemoveOrphanResults(t *testing.T) { + t.Run("Message", testPatchToolCallsRemoveOrphanResults[*schema.Message]) + t.Run("AgenticMessage", testPatchToolCallsRemoveOrphanResults[*schema.AgenticMessage]) +} + +func testPatchToolCallsRemoveDuplicateResults[M adk.MessageType](t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{RemoveDuplicateResults: true}) + require.NoError(t, err) + + state := &adk.TypedChatModelAgentState[M]{Messages: []M{ + makeAssistantMsgWithToolCalls[M]("", []testToolCall{{ID: "call_1", Name: "tool_a", Arguments: "{}"}}), + makeToolResultMsg[M]("result", "call_1", "tool_a"), + makeToolResultMsg[M]("duplicate", "call_1", "tool_a"), + }} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + assert.Equal(t, []string{"call_1"}, collectToolResultIDs(newState.Messages)) +} + +func TestPatchToolCallsRemoveDuplicateResults(t *testing.T) { + t.Run("Message", testPatchToolCallsRemoveDuplicateResults[*schema.Message]) + t.Run("AgenticMessage", testPatchToolCallsRemoveDuplicateResults[*schema.AgenticMessage]) +} + +func testPatchToolCallsSkipsEmptyIDInNonStrictMode[M adk.MessageType](t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[M](ctx, nil) + require.NoError(t, err) + + state := &adk.TypedChatModelAgentState[M]{Messages: []M{ + makeAssistantMsgWithToolCalls[M]("", []testToolCall{{ID: "", Name: "tool_a", Arguments: "{}"}}), + }} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + assert.Len(t, newState.Messages, 1) + assert.Empty(t, collectToolResultIDs(newState.Messages)) +} + +func TestPatchToolCallsSkipsEmptyIDInNonStrictMode(t *testing.T) { + t.Run("Message", testPatchToolCallsSkipsEmptyIDInNonStrictMode[*schema.Message]) + t.Run("AgenticMessage", testPatchToolCallsSkipsEmptyIDInNonStrictMode[*schema.AgenticMessage]) +} + +func testPatchToolCallsReportsEmptyIDInStrictMode[M adk.MessageType](t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{Strict: true}) + require.NoError(t, err) + + messages := []M{ + makeAssistantMsgWithToolCalls[M]("", []testToolCall{{ID: "", Name: "tool_a", Arguments: "{}"}}), + } + state := &adk.TypedChatModelAgentState[M]{Messages: messages} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.Error(t, err) + assert.Nil(t, newState) + assert.Same(t, any(messages[0]), any(state.Messages[0])) + assert.Contains(t, err.Error(), "empty_tool_call_id=1") +} + +func TestPatchToolCallsReportsEmptyIDInStrictMode(t *testing.T) { + t.Run("Message", testPatchToolCallsReportsEmptyIDInStrictMode[*schema.Message]) + t.Run("AgenticMessage", testPatchToolCallsReportsEmptyIDInStrictMode[*schema.AgenticMessage]) +} + +func TestPatchToolCallsStrictCountsAllMismatchCategories(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[*schema.Message](ctx, &Config{Strict: true}) + require.NoError(t, err) + + messages := []*schema.Message{ + makeToolResultMsg[*schema.Message]("orphan", "call_orphan", "tool_orphan"), + makeAssistantMsgWithToolCalls[*schema.Message]("", []testToolCall{ + {ID: "call_missing", Name: "tool_missing", Arguments: "{}"}, + {ID: "", Name: "tool_empty", Arguments: "{}"}, + {ID: "call_dup", Name: "tool_dup", Arguments: "{}"}, + }), + makeToolResultMsg[*schema.Message]("result", "call_dup", "tool_dup"), + makeToolResultMsg[*schema.Message]("duplicate", "call_dup", "tool_dup"), + } + state := &adk.TypedChatModelAgentState[*schema.Message]{Messages: messages} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.Error(t, err) + assert.Nil(t, newState) + assert.Equal(t, messages, state.Messages) + assert.Contains(t, err.Error(), "missing=1") + assert.Contains(t, err.Error(), "orphan=1") + assert.Contains(t, err.Error(), "duplicate=1") + assert.Contains(t, err.Error(), "empty_tool_call_id=1") +} + +func TestPatchToolCallsMarksSyntheticAgenticResult(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[*schema.AgenticMessage](ctx, &Config{MarkSynthetic: true}) + require.NoError(t, err) + + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: []*schema.AgenticMessage{ + makeAssistantMsgWithToolCalls[*schema.AgenticMessage]("", []testToolCall{{ID: "call_1", Name: "tool_a", Arguments: "{}"}}), + }} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + require.Len(t, newState.Messages, 2) + assertSyntheticMarker(t, newState.Messages[1], true) +} + +func TestPatchToolCallsMixedAgenticBlockRemovalUpdatesMessage(t *testing.T) { + ctx := context.Background() + assistant := makeAssistantMsgWithToolCalls[*schema.AgenticMessage]("", []testToolCall{{ID: "call_1", Name: "tool_a", Arguments: "{}"}}) + mixed := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.UserInputText{Text: "keep"}), + schema.NewContentBlock(&schema.FunctionToolResult{CallID: "call_orphan", Name: "tool_orphan"}), + schema.NewContentBlock(&schema.FunctionToolResult{CallID: "call_1", Name: "tool_a"}), + }, + } + adk.EnsureMessageID(mixed) + originalID := adk.GetMessageID(mixed) + + plan, err := buildAgenticNormalizationPlan(ctx, Config{RemoveOrphanResults: true}, []*schema.AgenticMessage{assistant, mixed}) + require.NoError(t, err) + require.Len(t, plan.messages, 2) + require.Len(t, plan.messages[1].ContentBlocks, 2) + assert.Equal(t, schema.ContentBlockTypeUserInputText, plan.messages[1].ContentBlocks[0].Type) + assert.Equal(t, "call_1", plan.messages[1].ContentBlocks[1].FunctionToolResult.CallID) + require.Len(t, plan.events, 1) + assert.Equal(t, adk.SessionEventMessageUpdated, plan.events[0].Kind) + assert.Equal(t, originalID, plan.events[0].MessageUpdated.MessageID) + assert.Equal(t, originalID, adk.GetMessageID(plan.events[0].MessageUpdated.Message)) +} + +func TestPatchToolCallsInsertionEventAnchorsReplayOrder(t *testing.T) { + ctx := context.Background() + assistant := makeAssistantMsgWithToolCalls[*schema.Message]("", []testToolCall{ + {ID: "call_1", Name: "tool_a", Arguments: "{}"}, + {ID: "call_2", Name: "tool_b", Arguments: "{}"}, + }) + result := makeToolResultMsg[*schema.Message]("result", "call_1", "tool_a") + messages := []*schema.Message{assistant, result} + + plan, err := buildMessageNormalizationPlan(ctx, Config{}, messages) + require.NoError(t, err) + require.Len(t, plan.messages, 3) + require.Len(t, plan.events, 1) + event := plan.events[0] + require.Equal(t, adk.SessionEventMessageInserted, event.Kind) + assert.Equal(t, adk.GetMessageID(result), event.MessageInserted.BeforeMessageID) + + replayed := append([]*schema.Message{}, messages...) + for i, msg := range replayed { + if adk.GetMessageID(msg) == event.MessageInserted.BeforeMessageID { + replayed = append(replayed, nil) + copy(replayed[i+1:], replayed[i:]) + replayed[i] = event.MessageInserted.Message + break + } + } + assert.Equal(t, []string{"call_2", "call_1"}, collectToolResultIDs(replayed)) + assert.Equal(t, []string{"call_2", "call_1"}, collectToolResultIDs(plan.messages)) +} + func TestPatchToolCallsAgenticToolSearchResult(t *testing.T) { ctx := context.Background() mw, err := NewTyped[*schema.AgenticMessage](ctx, nil) diff --git a/adk/middlewares/permission/permission.go b/adk/middlewares/permission/permission.go new file mode 100644 index 000000000..68d25c2c2 --- /dev/null +++ b/adk/middlewares/permission/permission.go @@ -0,0 +1,472 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package permission provides a ChatModelAgentMiddleware that gates tool execution +// behind a user-defined permission check. +package permission + +import ( + "context" + "fmt" + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[*AskInfo]("_eino_adk_permission_ask_info") + schema.RegisterName[*AskState]("_eino_adk_permission_ask_state") + schema.RegisterName[*DecisionEvent]("_eino_adk_permission_decision_event") +} + +// GateDecision is the result of a pre-execution permission check. +type GateDecision string + +const ( + // GateAllow bypasses the permission UI and executes the tool call. + GateAllow GateDecision = "allow" + // GateDeny skips tool execution and uses Message as the denial reason + // formatted through formatDenyResult. + GateDeny GateDecision = "deny" + // GateAsk interrupts the agent run for external approval. + GateAsk GateDecision = "ask" +) + +const ( + // SessionEventPermissionDecision records a valid user resume decision for a + // previously interrupted permission ask. + SessionEventPermissionDecision adk.SessionEventKind = adk.SessionEventKind(adk.SessionEventExtensionPrefix + "permission.decision") +) + +// GateCheckResult determines how a tool call should proceed before execution. +type GateCheckResult struct { + Decision GateDecision + + // Message is used as the deny reason or approval prompt. + Message string + + // UpdatedInput replaces ToolArgument.Text when the tool is allowed. + // Non-empty values are treated as replacements for backward compatibility. + UpdatedInput string + // HasUpdatedInput allows UpdatedInput to intentionally replace arguments with + // an empty string. + HasUpdatedInput bool + + // Reason is optional user-defined metadata for logging or auditing. + Reason string +} + +// Checker evaluates whether a tool call should be gated before execution. +// +// Returning an error signals an infrastructure failure and aborts the agent loop. +// Permission rejections should return GateDeny instead. Remembered preferences +// such as "always allow this action" should return GateAllow. +type Checker func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) + +// AskInfo is the user-facing interrupt payload emitted for Ask decisions. +type AskInfo struct { + ToolName string + Summary string `json:",omitempty"` +} + +// AskState is the private persisted interrupt state used to resume Ask decisions. +type AskState struct { + Info *AskInfo + + ToolName string + CallID string + Arguments string +} + +// ResumeAction resolves a previously interrupted permission ask. +type ResumeAction string + +const ( + // ResumeActionApprove executes the pending tool call. + ResumeActionApprove ResumeAction = "approve" + // ResumeActionReject rejects the pending tool call without execution. + ResumeActionReject ResumeAction = "reject" + // ResumeActionRespond returns alternate model-visible text without executing the tool. + ResumeActionRespond ResumeAction = "respond" +) + +// ResumeResponse is the data expected when resuming an Ask interrupt. +type ResumeResponse struct { + Action ResumeAction + + // UpdatedInput replaces the original arguments when Action is ResumeActionApprove. + // Non-empty values are treated as replacements for backward compatibility. + UpdatedInput string + // HasUpdatedInput allows UpdatedInput to intentionally replace arguments with + // an empty string. + HasUpdatedInput bool + + // Message is used as the rejection reason or model-visible response text. + Message string +} + +// DecisionEvent is the typed payload for SessionEventPermissionDecision. +// It intentionally omits the original saved tool arguments; only user-provided +// UpdatedInput is carried when it is part of an approval decision. +type DecisionEvent struct { + Action ResumeAction `json:"action"` + ToolName string `json:"tool_name"` + ToolUseID string `json:"tool_use_id,omitempty"` + DecisionText string `json:"decision_text,omitempty"` + UpdatedInput string `json:"updated_input,omitempty"` + HasUpdatedInput bool `json:"has_updated_input,omitempty"` +} + +// Middleware gates tool calls with a permission Checker. +type Middleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + checker Checker +} + +// NewTyped creates a typed permission middleware. +func NewTyped[M adk.MessageType](checker Checker) *Middleware[M] { + return &Middleware[M]{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[M]{}, + checker: checker, + } +} + +// New creates a permission middleware for the default *schema.Message agent path. +func New(checker Checker) *Middleware[*schema.Message] { + return NewTyped[*schema.Message](checker) +} + +type gateResult struct { + allowed bool + denyResult string + argument *schema.ToolArgument +} + +type normalizedResumeDecision struct { + Action ResumeAction + UpdatedInput string + HasUpdatedInput bool + DecisionText string +} + +func (m *Middleware[M]) permissionGate( + ctx context.Context, + tCtx *adk.ToolContext, + argument *schema.ToolArgument, +) (*gateResult, error) { + if argument == nil { + argument = &schema.ToolArgument{} + } + + wasInterrupted, hasState, savedState := tool.GetInterruptState[*AskState](ctx) + isTarget, hasData, response := tool.GetResumeContext[*ResumeResponse](ctx) + + if wasInterrupted && !hasState { + return &gateResult{allowed: true, argument: argument}, nil + } + + if wasInterrupted && !isTarget { + if !hasState || savedState == nil { + return nil, fmt.Errorf("permission: missing AskState for resumed tool %q (call_id=%s)", tCtx.Name, tCtx.CallID) + } + return nil, tool.StatefulInterrupt(ctx, savedState.publicInfo(), savedState) + } + + if isTarget && hasData { + if !hasState || savedState == nil { + return nil, fmt.Errorf("permission: missing AskState for targeted resume of tool %q (call_id=%s)", tCtx.Name, tCtx.CallID) + } + if err := emitDecisionEvent[M](ctx, tCtx, savedState, response); err != nil { + return nil, err + } + return handleResumeResponse(ctx, tCtx, &schema.ToolArgument{Text: savedState.Arguments}, response) + } + + if isTarget && !hasData { + return nil, fmt.Errorf( + "permission: tool %q (call_id=%s) was targeted for resume but received nil "+ + "or type-mismatched ResumeResponse; the caller must supply a *permission.ResumeResponse "+ + "via ResumeWithParams", tCtx.Name, tCtx.CallID) + } + + if m.checker == nil { + return nil, fmt.Errorf("permission: checker is nil for tool %q (call_id=%s)", tCtx.Name, tCtx.CallID) + } + + decision, err := m.checker(ctx, tCtx, argument) + if err != nil { + return nil, fmt.Errorf( + "permission: checker error for tool %q (call_id=%s, args=%s): %w", + tCtx.Name, tCtx.CallID, argument.Text, err) + } + if decision == nil { + return nil, fmt.Errorf( + "permission: checker returned nil GateCheckResult for tool %q (call_id=%s); "+ + "return a valid *GateCheckResult with Decision set to GateAllow, GateDeny, or GateAsk", + tCtx.Name, tCtx.CallID) + } + + switch decision.Decision { + case GateAllow: + return &gateResult{ + allowed: true, + argument: withUpdatedInput(argument, decision.UpdatedInput, decision.HasUpdatedInput || decision.UpdatedInput != ""), + }, nil + case GateDeny: + return &gateResult{denyResult: formatDenyResult(tCtx.Name, decision.Message)}, nil + case GateAsk: + info := &AskInfo{ + ToolName: tCtx.Name, + Summary: publicSummary(decision.Message, tCtx.CallID, argument.Text), + } + state := &AskState{ + Info: info, + ToolName: tCtx.Name, + CallID: tCtx.CallID, + Arguments: argument.Text, + } + return nil, tool.StatefulInterrupt(ctx, info, state) + case "": + return nil, fmt.Errorf("permission: empty gate decision for tool %q (call_id=%s); expected allow, deny, or ask", + tCtx.Name, tCtx.CallID) + default: + return nil, fmt.Errorf("permission: unknown gate decision %q for tool %q (call_id=%s); expected allow, deny, or ask", + decision.Decision, tCtx.Name, tCtx.CallID) + } +} + +func (s *AskState) publicInfo() *AskInfo { + if s == nil { + return nil + } + if s.Info != nil { + return s.Info + } + return &AskInfo{ToolName: s.ToolName} +} + +func publicSummary(message, callID, arguments string) string { + if message == "" { + return "" + } + if callID != "" && strings.Contains(message, callID) { + return "" + } + if arguments != "" && strings.Contains(message, arguments) { + return "" + } + return message +} + +func handleResumeResponse( + ctx context.Context, + tCtx *adk.ToolContext, + argument *schema.ToolArgument, + response *ResumeResponse, +) (*gateResult, error) { + decision, err := normalizeResumeDecision(tCtx, response) + if err != nil { + return nil, err + } + + switch decision.Action { + case ResumeActionApprove: + return &gateResult{ + allowed: true, + argument: withUpdatedInput(argument, decision.UpdatedInput, decision.HasUpdatedInput), + }, nil + case ResumeActionReject: + return &gateResult{denyResult: formatDenyResult(tCtx.Name, decision.DecisionText)}, nil + case ResumeActionRespond: + return &gateResult{denyResult: formatRespondResult(tCtx.Name, decision.DecisionText)}, nil + default: + return nil, fmt.Errorf("permission: unknown resume action %q for tool %q (call_id=%s); expected approve, reject, or respond", + decision.Action, tCtx.Name, tCtx.CallID) + } +} + +func normalizeResumeDecision(tCtx *adk.ToolContext, response *ResumeResponse) (*normalizedResumeDecision, error) { + toolName, callID := "", "" + if tCtx != nil { + toolName = tCtx.Name + callID = tCtx.CallID + } + if response == nil { + return nil, fmt.Errorf("permission: nil ResumeResponse for tool %q (call_id=%s)", toolName, callID) + } + + decision := &normalizedResumeDecision{Action: response.Action} + switch response.Action { + case ResumeActionApprove: + decision.HasUpdatedInput = response.HasUpdatedInput || response.UpdatedInput != "" + if decision.HasUpdatedInput { + decision.UpdatedInput = response.UpdatedInput + } + return decision, nil + case ResumeActionReject: + decision.DecisionText = response.Message + if decision.DecisionText == "" { + decision.DecisionText = "rejected by user" + } + return decision, nil + case ResumeActionRespond: + if response.Message == "" { + return nil, fmt.Errorf("permission: empty response message for respond action on tool %q (call_id=%s)", + toolName, callID) + } + decision.DecisionText = response.Message + return decision, nil + case "": + return nil, fmt.Errorf("permission: empty resume action for tool %q (call_id=%s); expected approve, reject, or respond", + toolName, callID) + default: + return nil, fmt.Errorf("permission: unknown resume action %q for tool %q (call_id=%s); expected approve, reject, or respond", + response.Action, toolName, callID) + } +} + +func emitDecisionEvent[M adk.MessageType](ctx context.Context, tCtx *adk.ToolContext, state *AskState, response *ResumeResponse) error { + if tCtx == nil { + return fmt.Errorf("permission: nil ToolContext for resume decision event") + } + if state == nil { + return fmt.Errorf("permission: nil AskState for resume decision event on tool %q (call_id=%s)", tCtx.Name, tCtx.CallID) + } + decision, err := normalizeResumeDecision(tCtx, response) + if err != nil { + return err + } + payload := &DecisionEvent{ + Action: decision.Action, + ToolName: state.ToolName, + ToolUseID: state.CallID, + DecisionText: decision.DecisionText, + UpdatedInput: decision.UpdatedInput, + HasUpdatedInput: decision.HasUpdatedInput, + } + return adk.TypedSendEvent[M](ctx, &adk.TypedAgentEvent[M]{ + SessionEventVariant: &adk.SessionEventVariant[M]{ + Event: &adk.SessionEvent[M]{ + Kind: SessionEventPermissionDecision, + Extension: &adk.SessionExtensionEvent{Data: payload}, + }, + }, + }) +} + +func (m *Middleware[M]) WrapInvokableToolCall( + _ context.Context, + endpoint adk.InvokableToolCallEndpoint, + tCtx *adk.ToolContext, +) (adk.InvokableToolCallEndpoint, error) { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + result, err := m.permissionGate(ctx, tCtx, &schema.ToolArgument{Text: argumentsInJSON}) + if err != nil { + return "", err + } + if !result.allowed { + return result.denyResult, nil + } + return endpoint(ctx, result.argument.Text, opts...) + }, nil +} + +func (m *Middleware[M]) WrapStreamableToolCall( + _ context.Context, + endpoint adk.StreamableToolCallEndpoint, + tCtx *adk.ToolContext, +) (adk.StreamableToolCallEndpoint, error) { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + result, err := m.permissionGate(ctx, tCtx, &schema.ToolArgument{Text: argumentsInJSON}) + if err != nil { + return nil, err + } + if !result.allowed { + return schema.StreamReaderFromArray([]string{result.denyResult}), nil + } + return endpoint(ctx, result.argument.Text, opts...) + }, nil +} + +func (m *Middleware[M]) WrapEnhancedInvokableToolCall( + _ context.Context, + endpoint adk.EnhancedInvokableToolCallEndpoint, + tCtx *adk.ToolContext, +) (adk.EnhancedInvokableToolCallEndpoint, error) { + return func(ctx context.Context, argument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + result, err := m.permissionGate(ctx, tCtx, argument) + if err != nil { + return nil, err + } + if !result.allowed { + return denyToolResult(result.denyResult), nil + } + return endpoint(ctx, result.argument, opts...) + }, nil +} + +func (m *Middleware[M]) WrapEnhancedStreamableToolCall( + _ context.Context, + endpoint adk.EnhancedStreamableToolCallEndpoint, + tCtx *adk.ToolContext, +) (adk.EnhancedStreamableToolCallEndpoint, error) { + return func(ctx context.Context, argument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + result, err := m.permissionGate(ctx, tCtx, argument) + if err != nil { + return nil, err + } + if !result.allowed { + return schema.StreamReaderFromArray([]*schema.ToolResult{denyToolResult(result.denyResult)}), nil + } + return endpoint(ctx, result.argument, opts...) + }, nil +} + +func withUpdatedInput(argument *schema.ToolArgument, updatedInput string, hasUpdatedInput bool) *schema.ToolArgument { + if !hasUpdatedInput { + return argument + } + cloned := *argument + cloned.Text = updatedInput + return &cloned +} + +func denyToolResult(denyMsg string) *schema.ToolResult { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: denyMsg}, + }, + } +} + +func formatDenyResult(toolName, message string) string { + tpl := internal.SelectPrompt(internal.I18nPrompts{ + English: "Permission denied for tool %s: %s", + Chinese: "工具 %s 权限被拒绝: %s", + }) + return fmt.Sprintf(tpl, toolName, message) +} + +func formatRespondResult(toolName, message string) string { + tpl := internal.SelectPrompt(internal.I18nPrompts{ + English: "Tool %s was not executed. User response: %s", + Chinese: "工具 %s 未执行。用户回复: %s", + }) + return fmt.Sprintf(tpl, toolName, message) +} diff --git a/adk/middlewares/permission/permission_test.go b/adk/middlewares/permission/permission_test.go new file mode 100644 index 000000000..77f58ff73 --- /dev/null +++ b/adk/middlewares/permission/permission_test.go @@ -0,0 +1,1372 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package permission + +import ( + "context" + "encoding/json" + "errors" + "io" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/internal/core" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +const addressSegmentAgent core.AddressSegmentType = "agent" + +func TestNewTypedSupportsBothMessageTypes(t *testing.T) { + checker := func(context.Context, *adk.ToolContext, *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAllow}, nil + } + + var _ adk.ChatModelAgentMiddleware = New(checker) + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = NewTyped[*schema.AgenticMessage](checker) +} + +func TestWrapInvokableToolCall_AllowWithUpdatedInput(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + assert.Equal(t, "WriteFile", tCtx.Name) + assert.Equal(t, "call_allow", tCtx.CallID) + assert.Equal(t, `{"path":"/etc/passwd"}`, args.Text) + return &GateCheckResult{Decision: GateAllow, UpdatedInput: `{"path":"/tmp/safe.txt"}`}, nil + }) + + var received string + endpoint := adk.InvokableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + received = argumentsInJSON + return "ok", nil + }) + + wrapped, err := m.WrapInvokableToolCall(context.Background(), endpoint, &adk.ToolContext{Name: "WriteFile", CallID: "call_allow"}) + require.NoError(t, err) + + result, err := wrapped(context.Background(), `{"path":"/etc/passwd"}`) + require.NoError(t, err) + assert.Equal(t, "ok", result) + assert.Equal(t, `{"path":"/tmp/safe.txt"}`, received) +} + +func TestWrapInvokableToolCall_AllowWithExplicitEmptyUpdatedInput(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAllow, HasUpdatedInput: true}, nil + }) + + received := "not called" + endpoint := adk.InvokableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + received = argumentsInJSON + return "ok", nil + }) + + wrapped, err := m.WrapInvokableToolCall(context.Background(), endpoint, &adk.ToolContext{Name: "WriteFile", CallID: "call_empty_update"}) + require.NoError(t, err) + + result, err := wrapped(context.Background(), `{"path":"/tmp/file"}`) + require.NoError(t, err) + assert.Equal(t, "ok", result) + assert.Empty(t, received) +} + +func TestWrapStreamableToolCall_Deny(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateDeny, Message: "blocked"}, nil + }) + + endpointCalled := false + endpoint := adk.StreamableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + endpointCalled = true + return schema.StreamReaderFromArray([]string{"unexpected"}), nil + }) + + wrapped, err := m.WrapStreamableToolCall(context.Background(), endpoint, &adk.ToolContext{Name: "Shell", CallID: "call_deny"}) + require.NoError(t, err) + + reader, err := wrapped(context.Background(), `{}`) + require.NoError(t, err) + require.NotNil(t, reader) + assert.False(t, endpointCalled) + + chunk, err := reader.Recv() + require.NoError(t, err) + assert.Equal(t, "Permission denied for tool Shell: blocked", chunk) + + _, err = reader.Recv() + assert.ErrorIs(t, err, io.EOF) +} + +func TestWrapInvokableToolCall_Respond(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve shell?"}, nil + }) + + endpointCalled := false + endpoint := adk.InvokableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + endpointCalled = true + return "unexpected", nil + }) + + tCtx := &adk.ToolContext{Name: "Shell", CallID: "call_standard_respond"} + wrapped, err := m.WrapInvokableToolCall(context.Background(), endpoint, tCtx) + require.NoError(t, err) + + _, err = wrapped(withAddress(context.Background()), `{"cmd":"rm -rf /"}`) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := wrapped(resumeContext(signal, &ResumeResponse{ + Action: ResumeActionRespond, + Message: "Explain first.", + }), `{"cmd":"rm -rf /"}`) + require.NoError(t, err) + assert.False(t, endpointCalled) + assert.Equal(t, formatRespondResult(tCtx.Name, "Explain first."), result) + assert.NotContains(t, result, "Permission denied") +} + +func TestWrapInvokableToolCall_ResumeApproveUsesSavedInterruptedArguments(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve write?"}, nil + }) + + var received string + endpoint := adk.InvokableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + received = argumentsInJSON + return "ok", nil + }) + + tCtx := &adk.ToolContext{Name: "WriteFile", CallID: "call_saved_args"} + wrapped, err := m.WrapInvokableToolCall(context.Background(), endpoint, tCtx) + require.NoError(t, err) + + _, err = wrapped(withAddress(context.Background()), `{"path":"/tmp/approved"}`) + require.Error(t, err) + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := wrapped(resumeContext(signal, &ResumeResponse{Action: ResumeActionApprove}), `{"path":"/etc/passwd"}`) + require.NoError(t, err) + assert.Equal(t, "ok", result) + assert.Equal(t, `{"path":"/tmp/approved"}`, received) +} + +func TestWrapInvokableToolCall_PassesThroughBusinessInterruptResume(t *testing.T) { + checkerCalls := 0 + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + checkerCalls++ + return &GateCheckResult{Decision: GateAsk, Message: "approve tool?"}, nil + }) + + endpointCalls := 0 + endpoint := adk.InvokableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + endpointCalls++ + wasInterrupted, hasState, state := tool.GetInterruptState[string](ctx) + isTarget, hasData, data := tool.GetResumeContext[string](ctx) + if wasInterrupted && hasState { + assert.Equal(t, "business-state", state) + require.True(t, isTarget) + require.True(t, hasData) + assert.Equal(t, "business-resume", data) + assert.Equal(t, `{"path":"/tmp/approved"}`, argumentsInJSON) + return "business resumed", nil + } + return "", tool.StatefulInterrupt(ctx, "business interrupt", "business-state") + }) + + tCtx := &adk.ToolContext{Name: "NestedTool", CallID: "call_nested_business"} + wrapped, err := m.WrapInvokableToolCall(context.Background(), endpoint, tCtx) + require.NoError(t, err) + + _, err = wrapped(withAddress(context.Background()), `{"path":"/tmp/approved"}`) + require.Error(t, err) + var permissionSignal *core.InterruptSignal + require.True(t, errors.As(err, &permissionSignal)) + + _, err = wrapped(resumeContext(permissionSignal, &ResumeResponse{Action: ResumeActionApprove}), `{"path":"/tmp/ignored"}`) + require.Error(t, err) + var businessSignal *core.InterruptSignal + require.True(t, errors.As(err, &businessSignal)) + + result, err := wrapped(genericResumeContext(businessSignal, "business-resume"), `{"path":"/tmp/approved"}`) + require.NoError(t, err) + assert.Equal(t, "business resumed", result) + assert.Equal(t, 1, checkerCalls) + assert.Equal(t, 2, endpointCalls) +} + +func TestAttack_BusinessInterruptNonTargetReplayPassesThrough(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAllow}, nil + }) + + endpointCalls := 0 + endpoint := adk.InvokableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + endpointCalls++ + wasInterrupted, hasState, state := tool.GetInterruptState[string](ctx) + isTarget, _, _ := tool.GetResumeContext[string](ctx) + if wasInterrupted { + require.True(t, hasState) + assert.Equal(t, "business-state", state) + require.False(t, isTarget) + return "", tool.StatefulInterrupt(ctx, "business interrupt", state) + } + return "", tool.StatefulInterrupt(ctx, "business interrupt", "business-state") + }) + + tCtx := &adk.ToolContext{Name: "NestedTool", CallID: "call_nested_business_nontarget"} + wrapped, err := m.WrapInvokableToolCall(context.Background(), endpoint, tCtx) + require.NoError(t, err) + + _, err = wrapped(withAddress(context.Background()), `{"path":"/tmp/approved"}`) + require.Error(t, err) + var businessSignal *core.InterruptSignal + require.True(t, errors.As(err, &businessSignal)) + + _, err = wrapped(nonTargetResumeContext(businessSignal), `{"path":"/tmp/approved"}`) + require.Error(t, err) + var replayedSignal *core.InterruptSignal + require.True(t, errors.As(err, &replayedSignal), "non-target replay should preserve the underlying business interrupt") + assert.NotContains(t, err.Error(), "missing AskState") + assert.Equal(t, 2, endpointCalls) +} + +func TestWrapInvokableToolCall_ResumeApproveWithExplicitEmptyUpdatedInput(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve empty override?"}, nil + }) + + received := "not called" + endpoint := adk.InvokableToolCallEndpoint(func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + received = argumentsInJSON + return "ok", nil + }) + + tCtx := &adk.ToolContext{Name: "WriteFile", CallID: "call_resume_empty_update"} + wrapped, err := m.WrapInvokableToolCall(context.Background(), endpoint, tCtx) + require.NoError(t, err) + + _, err = wrapped(withAddress(context.Background()), `{"path":"/tmp/approved"}`) + require.Error(t, err) + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := wrapped(resumeContext(signal, &ResumeResponse{Action: ResumeActionApprove, HasUpdatedInput: true}), `{"path":"/etc/passwd"}`) + require.NoError(t, err) + assert.Equal(t, "ok", result) + assert.Empty(t, received) +} + +func TestPermissionGate_AskThenResumeApprovedWithUpdatedInput(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve write?"}, nil + }) + + tCtx := &adk.ToolContext{Name: "WriteFile", CallID: "call_ask"} + ctx := withAddress(context.Background()) + + result, err := m.permissionGate(ctx, tCtx, &schema.ToolArgument{Text: `{"path":"/etc/passwd"}`}) + assert.Nil(t, result) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + require.NotNil(t, signal.InterruptState.State) + + askState, ok := signal.InterruptState.State.(*AskState) + require.True(t, ok) + require.NotNil(t, askState.Info) + assert.Equal(t, "WriteFile", askState.Info.ToolName) + assert.Equal(t, "call_ask", askState.CallID) + assert.Equal(t, `{"path":"/etc/passwd"}`, askState.Arguments) + + resumeCtx := resumeContext(signal, &ResumeResponse{ + Action: ResumeActionApprove, + UpdatedInput: `{"path":"/tmp/safe.txt"}`, + }) + + result, err = m.permissionGate(resumeCtx, tCtx, &schema.ToolArgument{Text: `{"path":"/etc/passwd"}`}) + require.NoError(t, err) + require.NotNil(t, result) + assert.True(t, result.allowed) + assert.Equal(t, `{"path":"/tmp/safe.txt"}`, result.argument.Text) +} + +func TestPermissionGate_AskPublicInfoOmitsPrivateFields(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: `approve call_public_info with {"path":"/etc/passwd"}?`}, nil + }) + + tCtx := &adk.ToolContext{Name: "WriteFile", CallID: "call_public_info"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{"path":"/etc/passwd"}`}) + require.Error(t, err) + + info := requireAskInfo(t, err) + assert.Equal(t, "WriteFile", info.ToolName) + + data, err := json.Marshal(info) + require.NoError(t, err) + got := string(data) + assert.Contains(t, got, "ToolName") + assert.NotContains(t, got, "CallID") + assert.NotContains(t, got, "Arguments") + assert.NotContains(t, got, "Message") + assert.NotContains(t, got, "call_public_info") + assert.NotContains(t, got, `{"path":"/etc/passwd"}`) +} + +func TestPermissionGate_AskPublicInfoIncludesSafeSummary(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "Approve running execute?"}, nil + }) + + tCtx := &adk.ToolContext{Name: "execute", CallID: "call_safe_summary"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{"cmd":"date"}`}) + require.Error(t, err) + + info := requireAskInfo(t, err) + assert.Equal(t, "execute", info.ToolName) + assert.Equal(t, "Approve running execute?", info.Summary) + + data, err := json.Marshal(info) + require.NoError(t, err) + got := string(data) + assert.Contains(t, got, "Summary") + assert.NotContains(t, got, "call_safe_summary") + assert.NotContains(t, got, `{"cmd":"date"}`) +} + +func TestPermissionGate_AskPublicInfoOmitsDuplicateSummary(t *testing.T) { + tests := []struct { + name string + message string + }{ + { + name: "call id", + message: "Approve call call_duplicate_summary?", + }, + { + name: "arguments", + message: `Approve running {"cmd":"rm -rf /"}?`, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: tt.message}, nil + }) + + tCtx := &adk.ToolContext{Name: "Shell", CallID: "call_duplicate_summary"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{"cmd":"rm -rf /"}`}) + require.Error(t, err) + + info := requireAskInfo(t, err) + assert.Empty(t, info.Summary) + + data, err := json.Marshal(info) + require.NoError(t, err) + got := string(data) + assert.NotContains(t, got, "Summary") + assert.NotContains(t, got, "call_duplicate_summary") + assert.NotContains(t, got, `{"cmd":"rm -rf /"}`) + }) + } +} + +func TestPermissionGate_ResumeApproveUsesAskStateArguments(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: `Approve call_private_args with {"path":"/tmp/approved"}?`}, nil + }) + + tCtx := &adk.ToolContext{Name: "WriteFile", CallID: "call_private_args"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{"path":"/tmp/approved"}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + askState, ok := signal.InterruptState.State.(*AskState) + require.True(t, ok) + require.NotNil(t, askState.Info) + require.Empty(t, askState.Info.Summary) + assert.Equal(t, `{"path":"/tmp/approved"}`, askState.Arguments) + + result, err := m.permissionGate(resumeContext(signal, &ResumeResponse{Action: ResumeActionApprove}), tCtx, &schema.ToolArgument{Text: `{"path":"/etc/passwd"}`}) + require.NoError(t, err) + require.NotNil(t, result) + assert.True(t, result.allowed) + assert.Equal(t, `{"path":"/tmp/approved"}`, result.argument.Text) +} + +func TestPermissionGate_AskThenResumeDenied(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve delete?"}, nil + }) + + tCtx := &adk.ToolContext{Name: "DeleteDB", CallID: "call_deny_resume"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := m.permissionGate(resumeContext(signal, &ResumeResponse{ + Action: ResumeActionReject, + Message: "user rejected", + }), tCtx, &schema.ToolArgument{Text: `{}`}) + require.NoError(t, err) + require.NotNil(t, result) + assert.False(t, result.allowed) + assert.Equal(t, "Permission denied for tool DeleteDB: user rejected", result.denyResult) +} + +func TestPermissionGate_AskThenResumeRespond(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve shell?"}, nil + }) + + tCtx := &adk.ToolContext{Name: "Shell", CallID: "call_respond"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{"cmd":"rm -rf /"}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := m.permissionGate(resumeContext(signal, &ResumeResponse{ + Action: ResumeActionRespond, + Message: "Please explain why this command is necessary first.", + }), tCtx, &schema.ToolArgument{Text: `{"cmd":"rm -rf /"}`}) + require.NoError(t, err) + require.NotNil(t, result) + assert.False(t, result.allowed) + assert.Equal(t, "Tool Shell was not executed. User response: Please explain why this command is necessary first.", result.denyResult) + assert.NotContains(t, result.denyResult, "Permission denied") +} + +func TestPermissionGate_ResumeRejectDoesNotExecute(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve delete?"}, nil + }) + + tCtx := &adk.ToolContext{Name: "DeleteDB", CallID: "call_reject_default"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := m.permissionGate(resumeContext(signal, &ResumeResponse{ + Action: ResumeActionReject, + }), tCtx, &schema.ToolArgument{Text: `{}`}) + require.NoError(t, err) + require.NotNil(t, result) + assert.False(t, result.allowed) + assert.Equal(t, "Permission denied for tool DeleteDB: rejected by user", result.denyResult) +} + +func TestPermissionGate_ResumeApproveWithUpdatedInput(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "sanitize?"}, nil + }) + + tCtx := &adk.ToolContext{Name: "WriteFile", CallID: "call_approve_update"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{"path":"/etc/passwd"}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := m.permissionGate(resumeContext(signal, &ResumeResponse{ + Action: ResumeActionApprove, + UpdatedInput: `{"path":"/tmp/safe.txt"}`, + }), tCtx, &schema.ToolArgument{Text: `{"path":"/etc/passwd"}`}) + require.NoError(t, err) + require.NotNil(t, result) + assert.True(t, result.allowed) + assert.Equal(t, `{"path":"/tmp/safe.txt"}`, result.argument.Text) +} + +func TestPermissionGate_InvalidResumeAction(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve?"}, nil + }) + + tCtx := &adk.ToolContext{Name: "Shell", CallID: "call_invalid_resume"} + _, err := m.permissionGate(withAddress(context.Background()), tCtx, &schema.ToolArgument{Text: `{}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := m.permissionGate(resumeContext(signal, &ResumeResponse{}), tCtx, &schema.ToolArgument{Text: `{}`}) + assert.Nil(t, result) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty resume action") + + result, err = m.permissionGate(resumeContext(signal, &ResumeResponse{Action: ResumeAction("unknown")}), tCtx, &schema.ToolArgument{Text: `{}`}) + assert.Nil(t, result) + require.Error(t, err) + assert.Contains(t, err.Error(), "unknown resume action") + + result, err = m.permissionGate(resumeContext(signal, &ResumeResponse{Action: ResumeActionRespond}), tCtx, &schema.ToolArgument{Text: `{}`}) + assert.Nil(t, result) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty response message") +} + +func TestPermissionGate_InvalidGateDecision(t *testing.T) { + tCtx := &adk.ToolContext{Name: "Shell", CallID: "call_invalid_gate"} + + tests := []struct { + name string + decision GateDecision + wantErr string + }{ + {name: "empty", decision: "", wantErr: "empty gate decision"}, + {name: "unknown", decision: GateDecision("unknown"), wantErr: "unknown gate decision"}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: tt.decision}, nil + }) + + result, err := m.permissionGate(context.Background(), tCtx, &schema.ToolArgument{Text: `{}`}) + assert.Nil(t, result) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestWrapEnhancedInvokableToolCall_Deny(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateDeny, Message: "enhanced blocked"}, nil + }) + + endpointCalled := false + endpoint := adk.EnhancedInvokableToolCallEndpoint(func(ctx context.Context, argument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + endpointCalled = true + return nil, nil + }) + + wrapped, err := m.WrapEnhancedInvokableToolCall(context.Background(), endpoint, &adk.ToolContext{Name: "Enhanced", CallID: "call_enhanced"}) + require.NoError(t, err) + + result, err := wrapped(context.Background(), &schema.ToolArgument{Text: `{}`}) + require.NoError(t, err) + assert.False(t, endpointCalled) + require.NotNil(t, result) + require.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "Permission denied for tool Enhanced: enhanced blocked", result.Parts[0].Text) +} + +func TestWrapEnhancedStreamableToolCall_AllowWithUpdatedInput(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAllow, UpdatedInput: `{"safe":true}`}, nil + }) + + var received string + endpoint := adk.EnhancedStreamableToolCallEndpoint(func(ctx context.Context, argument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + received = argument.Text + return schema.StreamReaderFromArray([]*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "ok"}}}, + }), nil + }) + + wrapped, err := m.WrapEnhancedStreamableToolCall(context.Background(), endpoint, &adk.ToolContext{Name: "EnhancedStream", CallID: "call_stream"}) + require.NoError(t, err) + + reader, err := wrapped(context.Background(), &schema.ToolArgument{Text: `{"unsafe":true}`}) + require.NoError(t, err) + require.NotNil(t, reader) + assert.Equal(t, `{"safe":true}`, received) + + chunk, err := reader.Recv() + require.NoError(t, err) + require.Len(t, chunk.Parts, 1) + assert.Equal(t, "ok", chunk.Parts[0].Text) +} + +func TestWrapEnhancedInvokableToolCall_Respond(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve enhanced?"}, nil + }) + + endpointCalled := false + endpoint := adk.EnhancedInvokableToolCallEndpoint(func(ctx context.Context, argument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + endpointCalled = true + return nil, nil + }) + + tCtx := &adk.ToolContext{Name: "Enhanced", CallID: "call_enhanced_respond"} + wrapped, err := m.WrapEnhancedInvokableToolCall(context.Background(), endpoint, tCtx) + require.NoError(t, err) + + _, err = wrapped(withAddress(context.Background()), &schema.ToolArgument{Text: `{}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + result, err := wrapped(resumeContext(signal, &ResumeResponse{ + Action: ResumeActionRespond, + Message: "Explain first.", + }), &schema.ToolArgument{Text: `{}`}) + require.NoError(t, err) + assert.False(t, endpointCalled) + require.NotNil(t, result) + require.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, formatRespondResult(tCtx.Name, "Explain first."), result.Parts[0].Text) +} + +func TestWrapEnhancedStreamableToolCall_Respond(t *testing.T) { + m := NewTyped[*schema.Message](func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve enhanced stream?"}, nil + }) + + endpointCalled := false + endpoint := adk.EnhancedStreamableToolCallEndpoint(func(ctx context.Context, argument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + endpointCalled = true + return nil, nil + }) + + tCtx := &adk.ToolContext{Name: "EnhancedStream", CallID: "call_enhanced_stream_respond"} + wrapped, err := m.WrapEnhancedStreamableToolCall(context.Background(), endpoint, tCtx) + require.NoError(t, err) + + _, err = wrapped(withAddress(context.Background()), &schema.ToolArgument{Text: `{}`}) + require.Error(t, err) + + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + + reader, err := wrapped(resumeContext(signal, &ResumeResponse{ + Action: ResumeActionRespond, + Message: "Use a safer approach.", + }), &schema.ToolArgument{Text: `{}`}) + require.NoError(t, err) + assert.False(t, endpointCalled) + require.NotNil(t, reader) + + chunk, err := reader.Recv() + require.NoError(t, err) + require.Len(t, chunk.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, chunk.Parts[0].Type) + assert.Equal(t, formatRespondResult(tCtx.Name, "Use a safer approach."), chunk.Parts[0].Text) + + _, err = reader.Recv() + assert.ErrorIs(t, err, io.EOF) +} + +func TestRespondFormattingIsByteIdenticalAcrossResultTypes(t *testing.T) { + want := formatRespondResult("ToolA", "continue without running") + assert.True(t, strings.HasPrefix(want, "Tool ToolA was not executed. User response: ")) + assert.Equal(t, want, denyToolResult(want).Parts[0].Text) +} + +func TestPermissionDecisionAppearsInToolUseTimeline(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + captureTool := &permissionCaptureTool{name: "permission_tool"} + info, err := captureTool.Info(ctx) + require.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "permission_call", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"path":"/tmp/file"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + checkerCalled := false + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "PermissionTimelineAgent", + Instruction: "use tools", + Model: cm, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + }, + }, + Handlers: []adk.ChatModelAgentMiddleware{ + New(func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + checkerCalled = true + return &GateCheckResult{Decision: GateAllow}, nil + }), + }, + }) + require.NoError(t, err) + + // In v3 the EvaluatedPermission field is removed from ToolSpanMeta. The gate + // decision is no longer surfaced on the tool span; for non-interrupted calls + // (gate=allow here), the decision is implicit in the tool result message + // content (a real tool invocation, not the deny prefix). We verify the + // real tool received its arguments and a tool_call_end span with status=ok + // was emitted. + var ( + sawToolCallEndOK bool + ) + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "permission-timeline", + SessionStore: &permissionSessionStore{}, + }) + iter := runner.Query(ctx, "use the tool", adk.WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant == nil || event.SessionEventVariant.Event == nil || event.SessionEventVariant.Event.Span == nil || event.SessionEventVariant.Event.Span.Tool == nil { + continue + } + if event.SessionEventVariant.Event.Kind == adk.SessionEventSpanToolCallEnd && event.SessionEventVariant.Event.Span.Status == "ok" { + sawToolCallEndOK = true + } + } + + assert.True(t, checkerCalled) + assert.True(t, sawToolCallEndOK, "expected a tool_call_end span with status=ok for the allow path") + assert.Equal(t, `{"path":"/tmp/file"}`, captureTool.received) +} + +func TestPermissionDecisionEventResumeLiveAndPersisted(t *testing.T) { + tests := []struct { + name string + response *ResumeResponse + wantAction ResumeAction + wantDecisionText string + wantUpdatedInput string + wantHasUpdated bool + wantToolInput string + wantToolNotInvoked bool + }{ + { + name: "approve with updated input", + response: &ResumeResponse{ + Action: ResumeActionApprove, + UpdatedInput: `{"path":"/tmp/safe.txt"}`, + }, + wantAction: ResumeActionApprove, + wantUpdatedInput: `{"path":"/tmp/safe.txt"}`, + wantHasUpdated: true, + wantToolInput: `{"path":"/tmp/safe.txt"}`, + }, + { + name: "approve with explicit empty updated input", + response: &ResumeResponse{Action: ResumeActionApprove, HasUpdatedInput: true}, + wantAction: ResumeActionApprove, + wantHasUpdated: true, + wantToolInput: "", + wantUpdatedInput: "", + }, + { + name: "reject with default text", + response: &ResumeResponse{Action: ResumeActionReject}, + wantAction: ResumeActionReject, + wantDecisionText: "rejected by user", + wantToolNotInvoked: true, + }, + { + name: "respond with decision text", + response: &ResumeResponse{ + Action: ResumeActionRespond, + Message: "Please explain first.", + }, + wantAction: ResumeActionRespond, + wantDecisionText: "Please explain first.", + wantToolNotInvoked: true, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + captureTool := &permissionCaptureTool{name: "permission_tool"} + info, err := captureTool.Info(ctx) + require.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "permission_call", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"path":"/etc/passwd"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "PermissionDecisionAgent", + Instruction: "use tools", + Model: cm, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + }, + }, + Handlers: []adk.ChatModelAgentMiddleware{ + New(func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: `Approve permission_call with {"path":"/etc/passwd"}?`}, nil + }), + }, + }) + require.NoError(t, err) + + sessionStore := &permissionSessionStore{} + checkpointStore := newPermissionCheckpointStore() + checkpointID := "permission-decision-" + strings.ReplaceAll(tt.name, " ", "-") + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + CheckPointStore: checkpointStore, + SessionID: checkpointID, + SessionStore: sessionStore, + }) + + var interruptID string + iter := runner.Query(ctx, "use the tool", adk.WithCheckPointID(checkpointID), adk.WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant == nil || event.SessionEventVariant.Event == nil || event.SessionEventVariant.Event.Kind != adk.SessionEventInterrupt { + continue + } + require.NotNil(t, event.SessionEventVariant.Event.Interrupt) + require.Len(t, event.SessionEventVariant.Event.Interrupt.Contexts, 1) + interruptID = event.SessionEventVariant.Event.Interrupt.Contexts[0].InterruptID + } + require.NotEmpty(t, interruptID) + + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{ + Targets: map[string]any{interruptID: tt.response}, + }, adk.WithTimelineEvents()) + require.NoError(t, err) + + var liveDecision *adk.SessionEvent[*schema.Message] + for { + event, ok := resumeIter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == SessionEventPermissionDecision { + liveDecision = event.SessionEventVariant.Event + } + } + requireDecisionEvent(t, liveDecision, tt.wantAction, tt.wantDecisionText, tt.wantUpdatedInput, tt.wantHasUpdated) + + decisions := filterPermissionDecisionEvents(sessionStore.events) + require.Len(t, decisions, 1) + requireDecisionEvent(t, decisions[0], tt.wantAction, tt.wantDecisionText, tt.wantUpdatedInput, tt.wantHasUpdated) + assert.Equal(t, liveDecision.EventID, decisions[0].EventID) + + decisionJSON, err := json.Marshal(decisions[0].Extension.Data) + require.NoError(t, err) + assert.NotContains(t, string(decisionJSON), `{"path":"/etc/passwd"}`) + assert.NotContains(t, string(decisionJSON), "Arguments") + assert.NotContains(t, string(decisionJSON), "CallID") + + decisionIndex, idleAfterDecisionIndex := -1, -1 + for i, event := range sessionStore.events { + if event.Kind == SessionEventPermissionDecision { + decisionIndex = i + } + if decisionIndex >= 0 && i > decisionIndex && event.Kind == adk.SessionEventSessionStatusIdle { + idleAfterDecisionIndex = i + break + } + } + require.NotEqual(t, -1, decisionIndex) + require.NotEqual(t, -1, idleAfterDecisionIndex) + assert.Less(t, decisionIndex, idleAfterDecisionIndex) + + if tt.wantToolNotInvoked { + assert.Empty(t, captureTool.received) + } else { + assert.Equal(t, tt.wantToolInput, captureTool.received) + } + }) + } +} + +func TestAttack_InvalidRespondDoesNotPersistDecisionEvent(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + captureTool := &permissionCaptureTool{name: "permission_tool"} + info, err := captureTool.Info(ctx) + require.NoError(t, err) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "permission_call", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"path":"/etc/passwd"}`}}, + }), nil).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "PermissionInvalidRespondAgent", + Instruction: "use tools", + Model: cm, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + }, + }, + Handlers: []adk.ChatModelAgentMiddleware{ + New(func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: "approve?"}, nil + }), + }, + }) + require.NoError(t, err) + + sessionStore := &permissionSessionStore{} + checkpointStore := newPermissionCheckpointStore() + const checkpointID = "permission-invalid-respond" + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + CheckPointStore: checkpointStore, + SessionID: checkpointID, + SessionStore: sessionStore, + }) + + var interruptID string + iter := runner.Query(ctx, "use the tool", adk.WithCheckPointID(checkpointID), adk.WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == adk.SessionEventInterrupt { + require.NotNil(t, event.SessionEventVariant.Event.Interrupt) + require.Len(t, event.SessionEventVariant.Event.Interrupt.Contexts, 1) + interruptID = event.SessionEventVariant.Event.Interrupt.Contexts[0].InterruptID + } + } + require.NotEmpty(t, interruptID) + + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{ + Targets: map[string]any{interruptID: &ResumeResponse{Action: ResumeActionRespond}}, + }, adk.WithTimelineEvents()) + require.NoError(t, err) + + var resumeErr error + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Err != nil { + resumeErr = event.Err + continue + } + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil { + assert.NotEqual(t, SessionEventPermissionDecision, event.SessionEventVariant.Event.Kind) + } + } + require.Error(t, resumeErr) + assert.Contains(t, resumeErr.Error(), "empty response message") + assert.Empty(t, filterPermissionDecisionEvents(sessionStore.events)) + assert.Empty(t, captureTool.received) +} + +// TestToolSpan_PermissionDenyEmitsBothSpansOnSameRun verifies plan §4.5.1 #6: +// when the permission gate denies on first invocation (no interrupt), the +// tool wrapper emits a tool_call_start + tool_call_end pair on the SAME run. +// The end span carries Status="ok" with a populated ToolResultMessageEventID +// — the deny content is the tool result, not an error. +func TestToolSpan_PermissionDenyEmitsBothSpansOnSameRun(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + captureTool := &permissionCaptureTool{name: "denied_tool"} + info, err := captureTool.Info(ctx) + require.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling", []schema.ToolCall{ + {ID: "deny_call", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"path":"/etc/passwd"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "PermissionDenyAgent", + Instruction: "use tools", + Model: cm, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + }, + }, + Handlers: []adk.ChatModelAgentMiddleware{ + New(func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateDeny, Message: "blocked"}, nil + }), + }, + }) + require.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "permission-deny-span", + SessionStore: &permissionSessionStore{}, + }) + + var ( + startSpanID string + startEventID string + endSpan *adk.SessionEvent[*schema.Message] + startCount int + endCount int + ) + + iter := runner.Query(ctx, "use the tool", adk.WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant == nil || event.SessionEventVariant.Event == nil || event.SessionEventVariant.Event.Span == nil || event.SessionEventVariant.Event.Span.Tool == nil { + continue + } + switch event.SessionEventVariant.Event.Kind { + case adk.SessionEventSpanToolCallStart: + startCount++ + startSpanID = event.SessionEventVariant.Event.Span.SpanID + startEventID = event.SessionEventVariant.Event.EventID + case adk.SessionEventSpanToolCallEnd: + endCount++ + endSpan = event.SessionEventVariant.Event + } + } + + assert.Equal(t, 1, startCount, "expected exactly one tool_call_start span on the deny run") + assert.Equal(t, 1, endCount, "expected exactly one tool_call_end span on the deny run") + require.NotNil(t, endSpan) + assert.Equal(t, startSpanID, endSpan.Span.SpanID, "end span shares SpanID with start span on the same run") + assert.Equal(t, startEventID, endSpan.Span.Tool.ToolCallStartEventID, "end span links back to start via ToolCallStartEventID") + assert.Equal(t, "ok", endSpan.Span.Status, "deny path produces a tool result (not an error); end span status is ok") + assert.NotEmpty(t, endSpan.Span.Tool.ToolResultMessageEventID, "deny end span must carry the ToolResultMessageEventID") + // The real tool must NOT have been invoked when the gate denies. + assert.Empty(t, captureTool.received, "deny path must not invoke the underlying tool") +} + +func TestPermissionGate_PersistedAgentInterruptOmitsPrivateInfo(t *testing.T) { + tests := []struct { + name string + message string + wantSummary bool + }{ + { + name: "safe summary", + message: "Approve running permission_tool?", + wantSummary: true, + }, + { + name: "duplicate message", + message: `Approve permission_call with {"path":"/etc/passwd"}?`, + wantSummary: false, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + captureTool := &permissionCaptureTool{name: "permission_tool"} + info, err := captureTool.Info(ctx) + require.NoError(t, err) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "permission_call", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"path":"/etc/passwd"}`}}, + }), nil).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "PermissionInterruptAgent", + Instruction: "use tools", + Model: cm, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + }, + }, + Handlers: []adk.ChatModelAgentMiddleware{ + New(func(ctx context.Context, tCtx *adk.ToolContext, args *schema.ToolArgument) (*GateCheckResult, error) { + return &GateCheckResult{Decision: GateAsk, Message: tt.message}, nil + }), + }, + }) + require.NoError(t, err) + + store := &permissionSessionStore{} + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "permission-agent-interrupt-" + strings.ReplaceAll(tt.name, " ", "-"), + SessionStore: store, + }) + iter := runner.Query(ctx, "use the tool", adk.WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + } + + var interrupt *adk.SessionEvent[*schema.Message] + for _, event := range store.events { + if event.Kind != adk.SessionEventInterrupt { + continue + } + interrupt = event + break + } + require.NotNil(t, interrupt) + require.NotNil(t, interrupt.Interrupt) + require.Len(t, interrupt.Interrupt.Contexts, 1) + + ctx0 := interrupt.Interrupt.Contexts[0] + assert.Equal(t, "permission_call", ctx0.ToolUseID) + + infoJSON, err := json.Marshal(ctx0.Info) + require.NoError(t, err) + infoText := string(infoJSON) + assert.Contains(t, infoText, "ToolName") + assert.Contains(t, infoText, "permission_tool") + assert.NotContains(t, infoText, "CallID") + assert.NotContains(t, infoText, "Arguments") + assert.NotContains(t, infoText, "Message") + assert.NotContains(t, infoText, "permission_call") + assert.NotContains(t, infoText, `{"path":"/etc/passwd"}`) + if tt.wantSummary { + assert.Contains(t, infoText, "Summary") + assert.Contains(t, infoText, tt.message) + } else { + assert.NotContains(t, infoText, "Summary") + assert.NotContains(t, infoText, tt.message) + } + assert.Empty(t, captureTool.received, "ask path must interrupt before invoking the underlying tool") + }) + } +} + +type permissionCaptureTool struct { + name string + received string +} + +func (t *permissionCaptureTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "permission capture tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "path": {Type: schema.String, Desc: "path"}, + }), + }, nil +} + +func (t *permissionCaptureTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + t.received = argumentsInJSON + return "ok", nil +} + +type permissionSessionStore struct { + events []*adk.SessionEvent[*schema.Message] +} + +func (s *permissionSessionStore) AppendEvents(_ context.Context, _ string, events []*adk.SessionEvent[*schema.Message]) error { + s.events = append(s.events, events...) + return nil +} + +func (s *permissionSessionStore) LoadEvents(_ context.Context, _ string, req *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[*schema.Message], error) { + if req == nil { + req = &adk.LoadSessionEventsRequest{} + } + start, end, step := 0, len(s.events), 1 + if req.Reverse { + start, end, step = len(s.events)-1, -1, -1 + } + if req.After != "" { + for i, event := range s.events { + if event != nil && event.EventID == req.After { + if req.Reverse { + start = i - 1 + } else { + start = i + 1 + } + break + } + } + } + var out []*adk.SessionEvent[*schema.Message] + hasMore := false + for i := start; i != end && i >= 0 && i < len(s.events); i += step { + if req.Limit > 0 && len(out) >= req.Limit { + hasMore = true + break + } + out = append(out, s.events[i]) + } + next := "" + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &adk.LoadSessionEventsResult[*schema.Message]{Events: out, Next: next}, nil +} + +type permissionCheckpointStore struct { + data map[string][]byte +} + +func newPermissionCheckpointStore() *permissionCheckpointStore { + return &permissionCheckpointStore{data: make(map[string][]byte)} +} + +func (s *permissionCheckpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + data, ok := s.data[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), data...), true, nil +} + +func (s *permissionCheckpointStore) Set(_ context.Context, key string, data []byte) error { + s.data[key] = append([]byte(nil), data...) + return nil +} + +func filterPermissionDecisionEvents(events []*adk.SessionEvent[*schema.Message]) []*adk.SessionEvent[*schema.Message] { + var decisions []*adk.SessionEvent[*schema.Message] + for _, event := range events { + if event.Kind == SessionEventPermissionDecision { + decisions = append(decisions, event) + } + } + return decisions +} + +func requireDecisionEvent( + t *testing.T, + event *adk.SessionEvent[*schema.Message], + action ResumeAction, + decisionText string, + updatedInput string, + hasUpdatedInput bool, +) { + t.Helper() + require.NotNil(t, event) + require.NotEmpty(t, event.EventID) + require.NotNil(t, event.Extension) + payload, ok := event.Extension.Data.(*DecisionEvent) + require.True(t, ok) + assert.Equal(t, action, payload.Action) + assert.Equal(t, "permission_tool", payload.ToolName) + assert.Equal(t, "permission_call", payload.ToolUseID) + assert.Equal(t, decisionText, payload.DecisionText) + assert.Equal(t, updatedInput, payload.UpdatedInput) + assert.Equal(t, hasUpdatedInput, payload.HasUpdatedInput) +} + +func requireAskInfo(t *testing.T, err error) *AskInfo { + t.Helper() + var signal *core.InterruptSignal + require.True(t, errors.As(err, &signal)) + info, ok := signal.InterruptInfo.Info.(*AskInfo) + require.True(t, ok) + require.NotNil(t, info) + return info +} + +func withAddress(ctx context.Context) context.Context { + return core.AppendAddressSegment(ctx, addressSegmentAgent, "test-agent", "") +} + +func resumeContext(signal *core.InterruptSignal, response *ResumeResponse) context.Context { + return genericResumeContext(signal, response) +} + +func genericResumeContext(signal *core.InterruptSignal, response any) context.Context { + id2Addr, id2State := core.SignalToPersistenceMaps(signal) + ctx := context.Background() + ctx = core.PopulateInterruptState(ctx, id2Addr, id2State) + ctx = core.BatchResumeWithData(ctx, map[string]any{signal.ID: response}) + return withAddress(ctx) +} + +func nonTargetResumeContext(signal *core.InterruptSignal) context.Context { + id2Addr, id2State := core.SignalToPersistenceMaps(signal) + ctx := context.Background() + ctx = core.PopulateInterruptState(ctx, id2Addr, id2State) + return withAddress(ctx) +} diff --git a/adk/middlewares/plantask/plantask.go b/adk/middlewares/plantask/plantask.go index fb201bddb..69ac74e20 100644 --- a/adk/middlewares/plantask/plantask.go +++ b/adk/middlewares/plantask/plantask.go @@ -63,7 +63,7 @@ type typedMiddleware[M adk.MessageType] struct { baseDir string } -func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { +func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { if runCtx == nil { return ctx, runCtx, nil } diff --git a/adk/middlewares/plantask/plantask_test.go b/adk/middlewares/plantask/plantask_test.go index 2354e79fd..6a76a70ce 100644 --- a/adk/middlewares/plantask/plantask_test.go +++ b/adk/middlewares/plantask/plantask_test.go @@ -62,7 +62,7 @@ func TestMiddlewareBeforeAgent(t *testing.T) { assert.NoError(t, err) assert.Nil(t, runCtx) - runCtx = &adk.ChatModelAgentContext{ + runCtx = &adk.ChatModelAgentContext[*schema.Message]{ Tools: []tool.BaseTool{}, } ctx, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) diff --git a/adk/middlewares/reduction/reduction.go b/adk/middlewares/reduction/reduction.go index fdd9931ff..890fc3e33 100644 --- a/adk/middlewares/reduction/reduction.go +++ b/adk/middlewares/reduction/reduction.go @@ -22,6 +22,7 @@ import ( "fmt" "io" "path/filepath" + "reflect" "strings" "unicode/utf8" @@ -31,6 +32,7 @@ import ( "github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" ) @@ -97,15 +99,22 @@ type TypedConfig[M adk.MessageType] struct { // TokenCounter is used to count the number of tokens in the conversation messages. // It is used to determine when to trigger clearing based on token usage, and token usage after clearing. - // Required. + // Optional. If not provided, a default token counter will be used, which estimates tokens by counting 1 token per 4 characters. TokenCounter func(ctx context.Context, msg []M, tools []*schema.ToolInfo) (int64, error) // MaxTokensForClear is the maximum number of tokens allowed in the conversation before clearing is attempted. // Required. Default is 160000. MaxTokensForClear int64 - // ClearRetentionSuffixLimit is the number of most recent messages to retain without clearing. - // This ensures the model has some immediate context. + // ClearRetentionSuffixLimit is the number of most recent tool-call rounds to retain without clearing. + // A round consists of one assistant message (which may contain multiple tool calls) and its corresponding tool-result messages. + // This ensures the model has immediate context to process pending tool results. + // + // Example with ClearRetentionSuffixLimit = 2, the retained suffix looks like: + // Round 2: assistant [call_A, call_B] → result_A, result_B + // Round 1: assistant [call_C] → result_C + // (Round 1 is the most recent, closest to the end of the message list.) + // // Optional. Default is 1. ClearRetentionSuffixLimit int @@ -360,6 +369,10 @@ type typedToolReductionMiddleware[M adk.MessageType] struct { excludeClearTools map[string]struct{} } +type clearRewriteDelta[M adk.MessageType] struct { + events []*adk.SessionEvent[M] +} + // getDefaultTokenCounter returns a default token counter function that operates on []M. // For *schema.Message it delegates to defaultTokenCounter. // For *schema.AgenticMessage it uses a simple character-based estimation. @@ -637,6 +650,7 @@ func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx con if estimatedTokens < t.config.MaxTokensForClear { return ctx, state, nil } + state.Messages = ensureMessageIDsOnCopiedMessages(state.Messages) // calc range var ( @@ -667,9 +681,10 @@ func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx con editTarget []M clearAtLeastTokens = t.config.ClearAtLeastTokens offloadStash []*offloadStashItem + pendingEvents []*adk.SessionEvent[M] ) - editTarget, end, err = t.applyClearRewriteGeneric(ctx, state, start, end, clearAtLeastTokens) + editTarget, end, pendingEvents, err = t.applyClearRewriteGeneric(ctx, state, start, end, clearAtLeastTokens) if err != nil { return ctx, state, err } @@ -681,16 +696,14 @@ func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx con toolCallMsg := editTarget[toolCallMsgIndex] toolCalls := getToolCallsGeneric(toolCallMsg) if isAssistantMsg(toolCallMsg) && len(toolCalls) > 0 { - toolMsgIndex := toolCallMsgIndex for _, tc := range toolCalls { - toolMsgIndex++ - if toolMsgIndex >= end { - break - } - resultMsg := editTarget[toolMsgIndex] - if !isToolResultMsg(resultMsg) { // unexpected - break + // Find the corresponding tool-result message by callID + resultMsgIndex, found := findToolResultByCallID(editTarget, toolCallMsgIndex+1, toolCallMsgIndex+1+len(toolCalls), tc.CallID) + if !found { + continue // No corresponding tool result found, skip } + resultMsg := editTarget[resultMsgIndex] + if _, found := t.excludeClearTools[tc.Name]; found { continue } @@ -744,10 +757,32 @@ func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx con setToolCallArguments(toolCallMsg, tc.BlockIndex, offloadInfo.ToolArgument.Text) setToolResultContent(resultMsg, offloadInfo.ToolResult, fromContent) + + // Queue MessageUpdated for the tool-result message (content replaced). + // ClearAtLeastTokens may still abort the clear, so persistence events + // must be emitted only after that threshold is satisfied. + pendingEvents = append(pendingEvents, &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessageUpdated, + MessageUpdated: &adk.MessageUpdatedEvent[M]{ + MessageID: adk.GetMessageID(resultMsg), + Message: resultMsg, + }, + }) } // set dedup flag setMsgClearedFlagGeneric(toolCallMsg) + + // Queue MessageUpdated for the assistant tool-call message (arguments + // rewritten + cleared flag set). Reconstruction must see this so the + // cleared flag suppresses double-reduction. + pendingEvents = append(pendingEvents, &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessageUpdated, + MessageUpdated: &adk.MessageUpdatedEvent[M]{ + MessageID: adk.GetMessageID(toolCallMsg), + Message: toolCallMsg, + }, + }) } toolCallMsgIndex++ } @@ -773,6 +808,12 @@ func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx con } } + for _, event := range pendingEvents { + if err := sendClearRewriteSessionEvent(ctx, event); err != nil { + return ctx, state, err + } + } + state.Messages = editTarget // replace original state messages if t.config.ClearPostProcess != nil { @@ -783,10 +824,11 @@ func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx con } func (t *typedToolReductionMiddleware[M]) applyClearRewriteGeneric(ctx context.Context, state *adk.TypedChatModelAgentState[M], start, end int, clearAtLeastTokens int64) ( - []M, int, error) { + []M, int, []*adk.SessionEvent[M], error) { var ( editTarget []M needProcessPart []M + delta clearRewriteDelta[M] ) editTarget = append(editTarget, state.Messages[:start]...) @@ -827,15 +869,25 @@ func (t *typedToolReductionMiddleware[M]) applyClearRewriteGeneric(ctx context.C } else { toolResponseMessages = needProcessPart[trStart:trEnd] } + spanEnd := trEnd + if spanEnd > len(needProcessPart) { + spanEnd = len(needProcessPart) + } + originalMessages := needProcessPart[i:spanEnd] rewrittenMessages, rewriteErr := t.config.ClearMessageRewriter(ctx, msg, toolResponseMessages) if rewriteErr != nil { - return nil, 0, rewriteErr + return nil, 0, nil, rewriteErr } + events, rewriteErr := buildClearRewriteEvents(originalMessages, rewrittenMessages) + if rewriteErr != nil { + return nil, 0, nil, rewriteErr + } + delta.events = append(delta.events, events...) rewritten = append(rewritten, rewrittenMessages...) i = trEnd } else { // unexpected - return nil, 0, fmt.Errorf("[applyClearRewrite] unexpected message: %v", any(msg)) + return nil, 0, nil, fmt.Errorf("[applyClearRewrite] unexpected message: %v", any(msg)) } } editTarget = append(editTarget, rewritten...) @@ -846,7 +898,202 @@ func (t *typedToolReductionMiddleware[M]) applyClearRewriteGeneric(ctx context.C editTarget = append(editTarget, state.Messages[end:]...) } - return editTarget, end, nil + return editTarget, end, delta.events, nil +} + +func sendClearRewriteSessionEvent[M adk.MessageType](ctx context.Context, event *adk.SessionEvent[M]) error { + err := adk.TypedSendEvent(ctx, &adk.TypedAgentEvent[M]{ + SessionEventVariant: &adk.SessionEventVariant[M]{Event: event}, + }) + if err != nil && strings.Contains(err.Error(), "must be called within a ChatModelAgent Run() or Resume() execution context") { + return nil + } + return err +} + +func buildClearRewriteEvents[M adk.MessageType](originalMessages []M, rewrittenMessages []M) ([]*adk.SessionEvent[M], error) { + originalIDs, err := messageIDsForRewrite("original", originalMessages, false) + if err != nil { + return nil, err + } + if len(rewrittenMessages) == 0 { + return []*adk.SessionEvent[M]{{ + Kind: adk.SessionEventMessagesDeleted, + MessagesDeleted: &adk.MessagesDeletedEvent{ + MessageIDs: originalIDs, + }, + }}, nil + } + rewrittenIDs, err := messageIDsForRewrite("rewritten", rewrittenMessages, true) + if err != nil { + return nil, err + } + if sameStringSlice(originalIDs, rewrittenIDs) { + var events []*adk.SessionEvent[M] + for i, msg := range rewrittenMessages { + if reflect.DeepEqual(originalMessages[i], msg) { + continue + } + events = append(events, &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessageUpdated, + MessageUpdated: &adk.MessageUpdatedEvent[M]{ + MessageID: rewrittenIDs[i], + Message: msg, + }, + }) + } + return events, nil + } + + originalIDSet := make(map[string]struct{}, len(originalIDs)) + for _, id := range originalIDs { + originalIDSet[id] = struct{}{} + } + var events []*adk.SessionEvent[M] + anchorID := originalIDs[0] + for i, msg := range rewrittenMessages { + if _, conflicts := originalIDSet[rewrittenIDs[i]]; conflicts { + msg = cloneMessageWithFreshID(msg) + rewrittenMessages[i] = msg + rewrittenIDs[i] = adk.GetMessageID(msg) + } + events = append(events, &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessageInserted, + MessageInserted: &adk.MessageInsertedEvent[M]{ + Message: msg, + BeforeMessageID: anchorID, + }, + }) + } + if err := validateUniqueIDs("rewritten", rewrittenIDs); err != nil { + return nil, err + } + events = append(events, &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessagesDeleted, + MessagesDeleted: &adk.MessagesDeletedEvent{ + MessageIDs: originalIDs, + }, + }) + return events, nil +} + +func messageIDsForRewrite[M adk.MessageType](label string, messages []M, ensure bool) ([]string, error) { + ids := make([]string, len(messages)) + for i, msg := range messages { + if ensure { + adk.EnsureMessageID(msg) + } + id := adk.GetMessageID(msg) + if id == "" { + return nil, fmt.Errorf("clear rewrite: %s message at index %d has empty message ID", label, i) + } + ids[i] = id + } + if err := validateUniqueIDs(label, ids); err != nil { + return nil, err + } + return ids, nil +} + +func validateUniqueIDs(label string, ids []string) error { + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + if _, ok := seen[id]; ok { + return fmt.Errorf("clear rewrite: %s messages contain duplicate message ID %q", label, id) + } + seen[id] = struct{}{} + } + return nil +} + +func sameStringSlice(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func cloneMessageWithFreshID[M adk.MessageType](msg M) M { + cloned := copyMessagesGeneric([]M{msg})[0] + switch m := any(cloned).(type) { + case *schema.Message: + if m.Extra != nil { + delete(m.Extra, internal.EinoMsgIDKey) + } + case *schema.AgenticMessage: + if m.Extra != nil { + delete(m.Extra, internal.EinoMsgIDKey) + } + } + adk.EnsureMessageID(cloned) + return cloned +} + +func ensureMessageIDsOnCopiedMessages[M adk.MessageType](msgs []M) []M { + var copied []M + for i, msg := range msgs { + if adk.GetMessageID(msg) != "" { + continue + } + if copied == nil { + copied = copyMessagesGeneric(msgs) + } + adk.EnsureMessageID(copied[i]) + } + if copied != nil { + return copied + } + return msgs +} + +func agenticResultCallID(block *schema.ContentBlock) (string, bool) { + if block == nil { + return "", false + } + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + return block.FunctionToolResult.CallID, true + } + if block.Type == schema.ContentBlockTypeToolSearchResult && block.ToolSearchFunctionToolResult != nil { + return block.ToolSearchFunctionToolResult.CallID, true + } + return "", false +} + +func findToolResultByCallID[M adk.MessageType](messages []M, startIndex, endIndex int, callID string) (int, bool) { + for i := startIndex; i < endIndex; i++ { + msg := messages[i] + if isToolResultMsg(msg) { + if getToolResultCallID(msg) == callID { + return i, true + } + } else { + return -1, false + } + } + return -1, false +} + +func getToolResultCallID[M adk.MessageType](msg M) string { + switch m := any(msg).(type) { + case *schema.Message: + if m.Role == schema.Tool { + return m.ToolCallID + } + case *schema.AgenticMessage: + if m.Role == schema.AgenticRoleTypeUser { + for _, block := range m.ContentBlocks { + if callID, ok := agenticResultCallID(block); ok { + return callID + } + } + } + } + return "" } type offloadStashItem struct { diff --git a/adk/middlewares/reduction/reduction_generic_test.go b/adk/middlewares/reduction/reduction_generic_test.go index b02d12b76..c43fdfd26 100644 --- a/adk/middlewares/reduction/reduction_generic_test.go +++ b/adk/middlewares/reduction/reduction_generic_test.go @@ -296,6 +296,49 @@ func testHelperFunctions[M adk.MessageType](t *testing.T) { copiedTCs := getToolCallsGeneric(copied[0]) assert.Equal(t, `{"modified":"true"}`, copiedTCs[0].Arguments) }) + + t.Run("getToolResultCallID", func(t *testing.T) { + // Tool result message with matching callID + tr := makeToolResultMsgG[M]("result content", "call_123", "my_tool") + assert.Equal(t, "call_123", getToolResultCallID(tr)) + + // Non-tool-result message should return empty string + user := makeUserMsgG[M]("hello") + assert.Equal(t, "", getToolResultCallID(user)) + }) + + t.Run("findToolResultByCallID", func(t *testing.T) { + messages := []M{ + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_1", Name: "tool1", Arguments: `{}`}, + }), + makeToolResultMsgG[M]("result 1", "call_1", "tool1"), + makeToolResultMsgG[M]("result 2", "call_2", "tool2"), + } + + // Find existing tool result + idx, found := findToolResultByCallID(messages, 1, 3, "call_2") + assert.True(t, found) + assert.Equal(t, 2, idx) + + // Find non-existent tool result + idx, found = findToolResultByCallID(messages, 1, 3, "call_999") + assert.False(t, found) + assert.Equal(t, -1, idx) + + // Stop searching when encountering a non-tool-result message + messages2 := []M{ + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_3", Name: "tool3", Arguments: `{}`}, + }), + makeToolResultMsgG[M]("result 3", "call_3", "tool3"), + makeUserMsgG[M]("not a tool result"), + makeToolResultMsgG[M]("result 4", "call_4", "tool4"), + } + idx, found = findToolResultByCallID(messages2, 1, 4, "call_4") + assert.False(t, found) + assert.Equal(t, -1, idx) + }) } // --------------------------------------------------------------------------- @@ -621,8 +664,8 @@ func TestToolResultFromMsgGeneric_AgenticMessage(t *testing.T) { { Type: schema.ContentBlockTypeFunctionToolResult, FunctionToolResult: &schema.FunctionToolResult{ - CallID: "c1", - Name: "tool1", + CallID: "c1", + Name: "tool1", Content: nil, }, }, diff --git a/adk/middlewares/reduction/reduction_test.go b/adk/middlewares/reduction/reduction_test.go index c22e9ec71..495bc152c 100644 --- a/adk/middlewares/reduction/reduction_test.go +++ b/adk/middlewares/reduction/reduction_test.go @@ -29,8 +29,11 @@ import ( "github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/session" + "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/compose" "github.com/cloudwego/eino/schema" ) @@ -353,7 +356,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }), - schema.ToolMessage("Sunny", "call_123456789"), + schema.ToolMessage("Sunny", "call_987654321"), schema.AssistantMessage("", []schema.ToolCall{ { ID: "call_123456789", @@ -460,7 +463,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }), - schema.ToolMessage("Sunny", "call_123456789"), + schema.ToolMessage("Sunny", "call_987654321"), schema.AssistantMessage("", []schema.ToolCall{ { ID: "call_123456789", @@ -491,6 +494,74 @@ func TestReductionMiddlewareClear(t *testing.T) { assert.Equal(t, "[Old tool result content cleared]", s.Messages[3].Content) }) + t.Run("test clear with reordered tool results", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + TokenCounter: defaultTokenCounter, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 0, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("get weather for two cities"), + // Assistant message with two tool calls + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "two_call_one", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London"}`}, + }, + { + ID: "two_call_two", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "Paris"}`}, + }, + }), + // Intentionally put tool result two BEFORE tool result one + schema.ToolMessage("Cloudy in Paris", "two_call_two"), + schema.ToolMessage("Sunny in London", "two_call_one"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + + // Verify tool call arguments are preserved (not offloaded since offload is true) + assert.Equal(t, `{"location": "London"}`, s.Messages[2].ToolCalls[0].Function.Arguments) + assert.Equal(t, `{"location": "Paris"}`, s.Messages[2].ToolCalls[1].Function.Arguments) + + // Verify tool result two is correctly matched and offloaded (it's at index 3) + assert.Equal(t, "Tool result saved to: /tmp/clear/two_call_two\nUse read_file to view", s.Messages[3].Content) + fileContent2, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/two_call_two", + }) + assert.NoError(t, err) + fileContentStr2 := strings.TrimPrefix(strings.TrimSpace(fileContent2.Content), "1\t") + assert.Equal(t, "Cloudy in Paris", fileContentStr2) + + // Verify tool result one is correctly matched and offloaded (it's at index 4) + assert.Equal(t, "Tool result saved to: /tmp/clear/two_call_one\nUse read_file to view", s.Messages[4].Content) + fileContent1, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/two_call_one", + }) + assert.NoError(t, err) + fileContentStr1 := strings.TrimPrefix(strings.TrimSpace(fileContent1.Content), "1\t") + assert.Equal(t, "Sunny in London", fileContentStr1) + }) + t.Run("test clear", func(t *testing.T) { backend := filesystem.NewInMemoryBackend() handler := func(ctx context.Context, detail *ToolDetail) (*ClearResult, error) { @@ -547,7 +618,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }), - schema.ToolMessage("Sunny", "call_123456789"), + schema.ToolMessage("Sunny", "call_987654321"), schema.AssistantMessage("", []schema.ToolCall{ { ID: "call_123456789", @@ -621,7 +692,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }), - schema.ToolMessage("Sunny", "call_123456789"), + schema.ToolMessage("Sunny", "call_987654321"), schema.AssistantMessage("", []schema.ToolCall{ { ID: "call_123456789", @@ -640,7 +711,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }, s.Messages[2].ToolCalls) - assert.NotNil(t, msgs[2].Extra[msgClearedFlag]) + assert.NotNil(t, s.Messages[2].Extra[msgClearedFlag]) assert.Equal(t, []schema.ToolCall{ { ID: "call_123456789", @@ -675,7 +746,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }, s.Messages[2].ToolCalls) - assert.NotNil(t, msgs[2].Extra[msgClearedFlag]) + assert.NotNil(t, s.Messages[2].Extra[msgClearedFlag]) assert.Equal(t, []schema.ToolCall{ { ID: "call_123456789", @@ -683,7 +754,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }, s.Messages[4].ToolCalls) - assert.NotNil(t, msgs[4].Extra[msgClearedFlag]) + assert.NotNil(t, s.Messages[4].Extra[msgClearedFlag]) assert.Equal(t, "Tool result saved to: /tmp/clear/call_987654321\nUse read_file to view", s.Messages[3].Content) assert.Equal(t, "Tool result saved to: /tmp/clear/call_123456789\nUse read_file to view", s.Messages[5].Content) }) @@ -872,7 +943,7 @@ func TestReductionMiddlewareClear(t *testing.T) { Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, }, }), - schema.ToolMessage("Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny", "call_123456789"), + schema.ToolMessage("Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny", "call_987654321"), schema.AssistantMessage("", []schema.ToolCall{ { ID: "call_123456789", @@ -2761,3 +2832,290 @@ func TestNewTypedAgenticMessage(t *testing.T) { var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw } + +func TestBuildClearRewriteEvents(t *testing.T) { + assistant := schema.AssistantMessage("", []schema.ToolCall{ + {ID: "call_1", Type: "function", Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file":"a"}`}}, + }) + toolMsg := schema.ToolMessage("ok", "call_1") + adk.EnsureMessageID(assistant) + adk.EnsureMessageID(toolMsg) + original := []adk.Message{assistant, toolMsg} + + t.Run("deletion", func(t *testing.T) { + events, err := buildClearRewriteEvents(original, nil) + assert.NoError(t, err) + assert.Len(t, events, 1) + assert.Equal(t, adk.SessionEventMessagesDeleted, events[0].Kind) + assert.Equal(t, []string{adk.GetMessageID(assistant), adk.GetMessageID(toolMsg)}, events[0].MessagesDeleted.MessageIDs) + }) + + t.Run("replacement inserts before delete", func(t *testing.T) { + replacement := schema.UserMessage("done") + events, err := buildClearRewriteEvents(original, []adk.Message{replacement}) + assert.NoError(t, err) + assert.Len(t, events, 2) + assert.Equal(t, adk.SessionEventMessageInserted, events[0].Kind) + assert.Equal(t, adk.GetMessageID(assistant), events[0].MessageInserted.BeforeMessageID) + assert.NotEmpty(t, adk.GetMessageID(events[0].MessageInserted.Message)) + assert.Equal(t, adk.SessionEventMessagesDeleted, events[1].Kind) + }) + + t.Run("same id content rewrite emits update", func(t *testing.T) { + updatedAssistant := schema.AssistantMessage("cleared", nil) + updatedAssistant.Extra = map[string]any{"_eino_msg_id": adk.GetMessageID(assistant)} + updatedTool := schema.ToolMessage("[placeholder]", "call_1") + updatedTool.Extra = map[string]any{"_eino_msg_id": adk.GetMessageID(toolMsg)} + events, err := buildClearRewriteEvents(original, []adk.Message{updatedAssistant, updatedTool}) + assert.NoError(t, err) + assert.Len(t, events, 2) + assert.Equal(t, adk.SessionEventMessageUpdated, events[0].Kind) + assert.Equal(t, adk.GetMessageID(assistant), events[0].MessageUpdated.MessageID) + assert.Equal(t, adk.SessionEventMessageUpdated, events[1].Kind) + assert.Equal(t, adk.GetMessageID(toolMsg), events[1].MessageUpdated.MessageID) + }) + + t.Run("duplicate replacement id errors", func(t *testing.T) { + a := schema.UserMessage("a") + b := schema.UserMessage("b") + dupID := "duplicate-id" + a.Extra = map[string]any{"_eino_msg_id": dupID} + b.Extra = map[string]any{"_eino_msg_id": dupID} + _, err := buildClearRewriteEvents(original, []adk.Message{a, b}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "duplicate") + }) + + t.Run("agentic deletion", func(t *testing.T) { + agenticAssistant := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: &schema.FunctionToolCall{ + CallID: "agentic-call", + Name: "write_file", + }, + }, + }, + } + agenticTool := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "agentic-call", + Name: "write_file", + }, + }, + }, + } + adk.EnsureMessageID(agenticAssistant) + adk.EnsureMessageID(agenticTool) + events, err := buildClearRewriteEvents([]*schema.AgenticMessage{agenticAssistant, agenticTool}, nil) + assert.NoError(t, err) + assert.Len(t, events, 1) + assert.Equal(t, adk.SessionEventMessagesDeleted, events[0].Kind) + }) +} + +type reductionRewritePersistModel struct { + calls int + inputs [][]*schema.Message +} + +func (m *reductionRewritePersistModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.calls++ + m.inputs = append(m.inputs, copyMessages(input)) + if m.calls == 1 { + return schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "mock_invokable_tool", Arguments: `{"value":"x"}`}, + }, + }), nil + } + return schema.AssistantMessage("done", nil), nil +} + +func (m *reductionRewritePersistModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func TestClearMessageRewriterPersistsMessagesDeletedThroughRunner(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + model := &reductionRewritePersistModel{} + mw, err := New(ctx, &Config{ + SkipTruncation: true, + MaxTokensForClear: 1, + ClearRetentionSuffixLimit: -1, + TokenCounter: func(context.Context, []adk.Message, []*schema.ToolInfo) (int64, error) { + return 1000, nil + }, + ClearMessageRewriter: func(context.Context, adk.Message, []adk.Message) ([]adk.Message, error) { + return nil, nil + }, + }) + assert.NoError(t, err) + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "reduction-delete-agent", + Description: "reduction delete test agent", + Model: model, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{mockInvokableTool()}}}, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + assert.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "reduction-delete-session", + SessionStore: store, + }) + drainReductionEvents(t, runner.Query(ctx, "please call the tool")) + + events := loadReductionSessionEvents(t, ctx, store, "reduction-delete-session") + var deletedIDs []string + for _, event := range events { + if event.MessagesDeleted != nil { + deletedIDs = event.MessagesDeleted.MessageIDs + } + } + assert.Len(t, deletedIDs, 2) + + nextModel := &reductionRewritePersistModel{} + nextAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "reduction-delete-agent", + Description: "reduction delete test agent", + Model: nextModel, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{mockInvokableTool()}}}, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + assert.NoError(t, err) + nextRunner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: nextAgent, + SessionID: "reduction-delete-session", + SessionStore: store, + }) + drainReductionEvents(t, nextRunner.Query(ctx, "next turn")) + + if assert.NotEmpty(t, nextModel.inputs) { + for _, msg := range nextModel.inputs[0] { + assert.False(t, msg.Role == schema.Tool && msg.ToolCallID == "call_1") + for _, tc := range msg.ToolCalls { + assert.NotEqual(t, "call_1", tc.ID) + } + } + } +} + +func TestClearMessageRewriterAbortDoesNotPersistStructuralEvents(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + model := &reductionRewritePersistModel{} + callCount := 0 + mw, err := New(ctx, &Config{ + SkipTruncation: true, + MaxTokensForClear: 1, + ClearRetentionSuffixLimit: -1, + ClearAtLeastTokens: 10, + TokenCounter: func(context.Context, []adk.Message, []*schema.ToolInfo) (int64, error) { + callCount++ + if callCount == 1 { + return 1000, nil + } + return 999, nil + }, + ClearMessageRewriter: func(context.Context, adk.Message, []adk.Message) ([]adk.Message, error) { + return nil, nil + }, + }) + assert.NoError(t, err) + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "reduction-abort-agent", + Description: "reduction abort test agent", + Model: model, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{mockInvokableTool()}}}, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + assert.NoError(t, err) + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "reduction-abort-session", + SessionStore: store, + }) + drainReductionEvents(t, runner.Query(ctx, "please call the tool")) + + events := loadReductionSessionEvents(t, ctx, store, "reduction-abort-session") + for _, event := range events { + assert.Nil(t, event.MessageUpdated) + assert.Nil(t, event.MessageInserted) + assert.Nil(t, event.MessagesDeleted) + } +} + +func TestClearAtLeastTokensAbortDoesNotPersistMessageUpdates(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + backend := filesystem.NewInMemoryBackend() + model := &reductionRewritePersistModel{} + callCount := 0 + mw, err := New(ctx, &Config{ + Backend: backend, + SkipTruncation: true, + MaxTokensForClear: 1, + ClearRetentionSuffixLimit: -1, + ClearAtLeastTokens: 10, + TokenCounter: func(context.Context, []adk.Message, []*schema.ToolInfo) (int64, error) { + callCount++ + if callCount == 1 { + return 1000, nil + } + return 999, nil + }, + }) + assert.NoError(t, err) + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "reduction-clear-abort-agent", + Description: "reduction clear abort test agent", + Model: model, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{mockInvokableTool()}}}, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + assert.NoError(t, err) + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "reduction-clear-abort-session", + SessionStore: store, + }) + drainReductionEvents(t, runner.Query(ctx, "please call the tool")) + + events := loadReductionSessionEvents(t, ctx, store, "reduction-clear-abort-session") + for _, event := range events { + assert.Nil(t, event.MessageUpdated) + } +} + +func drainReductionEvents(t *testing.T, iter *adk.AsyncIterator[*adk.AgentEvent]) { + t.Helper() + for { + event, ok := iter.Next() + if !ok { + return + } + assert.NoError(t, event.Err) + } +} + +func loadReductionSessionEvents(t *testing.T, ctx context.Context, store adk.SessionEventStore[*schema.Message], sessionID string) []*adk.SessionEvent[*schema.Message] { + t.Helper() + res, err := store.LoadEvents(ctx, sessionID, &adk.LoadSessionEventsRequest{}) + assert.NoError(t, err) + return res.Events +} diff --git a/adk/middlewares/skill/skill.go b/adk/middlewares/skill/skill.go index 8f8b2cad3..940d71f84 100644 --- a/adk/middlewares/skill/skill.go +++ b/adk/middlewares/skill/skill.go @@ -272,7 +272,7 @@ type typedSkillHandler[M adk.MessageType] struct { tool *typedSkillTool[M] } -func (h *typedSkillHandler[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { +func (h *typedSkillHandler[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { runCtx.Instruction = runCtx.Instruction + "\n" + h.instruction runCtx.Tools = append(runCtx.Tools, h.tool) return ctx, runCtx, nil diff --git a/adk/middlewares/skill/skill_test.go b/adk/middlewares/skill/skill_test.go index 3cc536abd..5c0596bab 100644 --- a/adk/middlewares/skill/skill_test.go +++ b/adk/middlewares/skill/skill_test.go @@ -456,7 +456,7 @@ func TestBeforeAgent(t *testing.T) { handler, err := NewMiddleware(ctx, &Config{Backend: backend}) require.NoError(t, err) - runCtx := &adk.ChatModelAgentContext{ + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ Instruction: "base instruction", Tools: []tool.BaseTool{}, } diff --git a/adk/middlewares/subagent/agent_tool.go b/adk/middlewares/subagent/agent_tool.go new file mode 100644 index 000000000..6b3a392c5 --- /dev/null +++ b/adk/middlewares/subagent/agent_tool.go @@ -0,0 +1,219 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package subagent + +import ( + "context" + "fmt" + "path/filepath" + "strings" + + "github.com/bytedance/sonic" + "github.com/google/uuid" + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/compose" +) + +const ( + agentToolName = "agent" + // TaskTypeSubagent is the backgroundtask Task.Type tag for sub-agent tasks + // launched by the agent tool, letting a shared Manager distinguish them from + // shell tasks. + TaskTypeSubagent = "subagent" + + // MetadataKeySubagentType is the RunInput.Metadata / Task.Metadata key under + // which the agent tool records the sub-agent type for a task. A + // ShouldAutoBackground hook reads it (via TypeFromTask) to apply + // agent-type-specific policy without parsing the human-readable Description. The + // value is a string. + MetadataKeySubagentType = "subagent_type" +) + +// TypeFromTask returns the sub-agent type recorded in a sub-agent task's +// metadata under MetadataKeySubagentType, or "" if absent (e.g. the task is not a +// sub-agent run). It is the intended way for a ShouldAutoBackground hook to recover +// the agent type. +func TypeFromTask(t *backgroundtask.Task) string { + if t == nil { + return "" + } + st, _ := t.Metadata[MetadataKeySubagentType].(string) + return st +} + +// agentInput is the agent tool's input when no Manager is configured: spawn a +// sub-agent synchronously in the foreground. +type agentInput struct { + SubagentType string `json:"subagent_type" jsonschema:"required" jsonschema_description:"The type of specialized agent to use for this task"` + Prompt string `json:"prompt" jsonschema:"required" jsonschema_description:"The task for the agent to perform"` + Description string `json:"description" jsonschema:"required" jsonschema_description:"A short (3-5 word) description of the task"` +} + +// agentManagedInput is the agent tool's input when a Manager is configured: it adds +// run_in_background so the model can spawn the sub-agent in the background. +type agentManagedInput struct { + agentInput + RunInBackground bool `json:"run_in_background,omitempty" jsonschema_description:"Set to true to run this agent in the background. You will be notified when it completes."` +} + +// newAgentTool builds the foreground-only agent tool (no Manager): it invokes the +// agent-as-tool adapter directly, forwarding opts so event forwarding, session +// sharing and interrupt/resume behave exactly as a normal agent-as-tool call. +func newAgentTool(subAgents map[string]tool.InvokableTool, name, desc string) (tool.BaseTool, error) { + return utils.InferOptionableTool(name, desc, + func(ctx context.Context, in agentInput, opts ...tool.Option) (string, error) { + a, params, err := resolveSubAgent(subAgents, in.SubagentType, in.Prompt, in.Description) + if err != nil { + return "", err + } + return a.InvokableRun(ctx, params, opts...) + }) +} + +// newManagedAgentTool builds the Manager-backed agent tool. It wraps the same +// agent-as-tool invocation in a managed task, so foreground behavior is identical +// and only lifecycle/background switching is layered on top. +// +// When store and outputDir are both set, each run is given an output file at +// outputDir/.output: the file is created empty up front so its advertised +// path exists immediately, and the sub-agent's final result is appended there on +// completion. The Manager never writes — the tool owns it. store is a +// filesystem.Appender; output files require one (no rewrite fallback). +func newManagedAgentTool(mgr *backgroundtask.Manager, subAgents map[string]tool.InvokableTool, store filesystem.Appender, outputDir, name, desc string) (tool.BaseTool, error) { + return utils.InferOptionableTool(name, desc, + func(ctx context.Context, in agentManagedInput, opts ...tool.Option) (string, error) { + a, params, err := resolveSubAgent(subAgents, in.SubagentType, in.Prompt, in.Description) + if err != nil { + return "", err + } + + outputFile := reserveAgentOutputFile(ctx, store, outputDir) + + result, err := mgr.Run(ctx, &backgroundtask.RunInput{ + Description: in.Description, + Type: TaskTypeSubagent, + ToolUseID: compose.GetToolCallID(ctx), + RunInBackground: in.RunInBackground, + Metadata: map[string]any{MetadataKeySubagentType: in.SubagentType}, + OutputFile: outputFile, + }, func(workCtx context.Context, task backgroundtask.TaskInfo) (string, error) { + out, runErr := a.InvokableRun(workCtx, params, opts...) + if runErr != nil { + return "", runErr + } + if outputFile != "" { + if appendErr := store.Append(workCtx, &filesystem.AppendRequest{FilePath: outputFile, Content: out}); appendErr != nil { + // The result never reached the file: mark it unreliable (by task id) + // so task_output reports the file's failed state instead of trusting + // the empty/partial file. + mgr.MarkOutputFileUnreliable(task.ID, appendErr.Error()) + } + } + return out, nil + }) + if err != nil { + return "", err + } + + switch result.Status { + case backgroundtask.StatusCompleted: + return result.Result, nil + case backgroundtask.StatusRunning: + msg := fmt.Sprintf("Agent running in background with ID: %s.", result.ID) + if result.OutputFile != "" { + msg += fmt.Sprintf(" Output is being written to: %s.", result.OutputFile) + } + msg += " You will be notified when it completes." + if result.OutputFile != "" { + msg += " To check interim output, use Read on that file path." + } + return msg, nil + case backgroundtask.StatusFailed: + return "", fmt.Errorf("subagent %q task %q (%s) failed: %s", + in.SubagentType, result.ID, in.Description, result.Error) + case backgroundtask.StatusCanceled: + return "", fmt.Errorf("subagent %q task %q (%s) was canceled", + in.SubagentType, result.ID, in.Description) + default: + return result.Result, nil + } + }) +} + +// reserveAgentOutputFile reserves an output-file path under outputDir and creates +// it empty (via Append) so the path exists before the run completes. The file is +// named after the launching tool-call id (so it matches Task.ToolUseID), falling +// back to a uuid when no tool-call id is in context. Returns "" when output files +// are not configured (no store / no dir) or when the up-front reservation write +// fails — in the latter case the task advertises no output file, so consumers +// fall back to the in-memory Result. +func reserveAgentOutputFile(ctx context.Context, store filesystem.Appender, outputDir string) string { + if store == nil || outputDir == "" { + return "" + } + name := compose.GetToolCallID(ctx) + if name == "" { + name = uuid.NewString() + } + path := filepath.Join(outputDir, name+".output") + if err := store.Append(ctx, &filesystem.AppendRequest{FilePath: path, Content: ""}); err != nil { + return "" + } + return path +} + +// resolveSubAgent looks up the agent-as-tool adapter for subagentType and builds +// the marshaled request for it. If prompt is empty, description is used as the +// task request. +func resolveSubAgent(subAgents map[string]tool.InvokableTool, subagentType, prompt, description string) (tool.InvokableTool, string, error) { + a, ok := subAgents[subagentType] + if !ok { + return nil, "", fmt.Errorf("subagent type %q not found", subagentType) + } + if prompt == "" { + prompt = description + } + params, err := sonic.MarshalString(map[string]string{"request": prompt}) + if err != nil { + return nil, "", err + } + return a, params, nil +} + +// defaultAgentToolDescription generates the agent tool description with sub-agent list. +func defaultAgentToolDescription[M adk.MessageType](ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error) { + subAgentsDescBuilder := strings.Builder{} + for _, a := range subAgents { + name := a.Name(ctx) + desc := a.Description(ctx) + _, _ = fmt.Fprintf(&subAgentsDescBuilder, "- %s: %s\n", name, desc) + } + toolDesc := internal.SelectPrompt(internal.I18nPrompts{ + English: agentToolDescription, + Chinese: agentToolDescriptionChinese, + }) + return pyfmt.Fmt(toolDesc, map[string]any{ + "other_agents": subAgentsDescBuilder.String(), + }) +} diff --git a/adk/middlewares/subagent/middleware.go b/adk/middlewares/subagent/middleware.go new file mode 100644 index 000000000..d1580ce16 --- /dev/null +++ b/adk/middlewares/subagent/middleware.go @@ -0,0 +1,202 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package subagent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// Config configures the subagent middleware for the standard *schema.Message message type. +// It is the default specialization of TypedConfig. +type Config = TypedConfig[*schema.Message] + +// TypedConfig configures the subagent middleware, parameterized by message type. +type TypedConfig[M adk.MessageType] struct { + // SubAgents is the list of agents available for spawning. + // Each agent must have a unique name. Required. + SubAgents []adk.TypedAgent[M] + + // ToolName overrides the name of the agent-spawning tool. + // When empty, defaults to "agent". + ToolName string + + // ToolDescriptionGenerator overrides the default agent tool description generator. + // The generator receives the list of sub-agents and should return a complete tool + // description string. When nil, defaultAgentToolDescription is used. + ToolDescriptionGenerator func(ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error) + + // SystemPrompt overrides the default system prompt injected by BeforeAgent. + // When nil, the built-in prompt (with i18n support) is used. + // Defined as *string because an empty string may be an intentional user value. + SystemPrompt *string + + // Background configures background-task execution for sub-agent runs. When nil, + // only foreground (blocking) agent execution is available and runs are NOT + // tracked. See BackgroundConfig. + Background *BackgroundConfig +} + +// BackgroundConfig enables background-task execution for the agent tool. +// +// When set, ALL agent runs (foreground and background) are managed by the Manager, +// making them visible via Get/List, and the Agent tool gains a run_in_background +// parameter. +type BackgroundConfig struct { + // Manager is the shared background-task Manager. Required (a nil Manager is the + // same as no BackgroundConfig). It may be shared with other middlewares (e.g. + // filesystem) so a single task-ID space spans agent and shell runs. The + // task_output/task_stop control tools are NOT injected here; wire the + // backgroundtask control middleware (adk/middlewares/backgroundtask) once, bound + // to the same Manager. + Manager *backgroundtask.Manager + + // OutputStore and OutputDir, when both set, give every managed sub-agent run an + // output file at OutputDir/.output. The managed agent tool appends the + // sub-agent's final result there on completion and records the path on + // Task.OutputFile, so a backgrounded run's result is retrievable by path (and + // large results need not be inlined). The Manager itself never writes. + // OutputStore is a filesystem.Appender (filesystem.InMemoryBackend implements + // it); output files require one. When either is unset, runs have no output file. + OutputStore filesystem.Appender + OutputDir string +} + +// New creates a ChatModelAgentMiddleware that injects sub-agent tools into the agent context. +// +// The middleware injects an Agent tool for spawning sub-agents. When Config.Manager is +// provided, agent runs are tracked by the shared background-task Manager and the Agent +// tool gains a run_in_background parameter. The task_output/task_stop control tools are +// NOT injected here; wire the backgroundtask control middleware +// (adk/middlewares/backgroundtask) once, bound to the same Manager. +func New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped[*schema.Message](ctx, config) +} + +// NewTyped creates a TypedChatModelAgentMiddleware that injects sub-agent tools into the +// agent context, parameterized by message type. See New for behavior details. +func NewTyped[M adk.MessageType](ctx context.Context, config *TypedConfig[M]) (adk.TypedChatModelAgentMiddleware[M], error) { + if err := validate(ctx, config); err != nil { + return nil, err + } + + // Build subAgentToolMap: name → the agent-as-tool adapter that runs the agent. + // Both the foreground and the Manager-backed paths invoke this same adapter. + subAgentToolMap := make(map[string]tool.InvokableTool, len(config.SubAgents)) + for _, a := range config.SubAgents { + name := a.Name(ctx) + bt := adk.NewTypedAgentTool[M](ctx, a) + it, ok := bt.(tool.InvokableTool) + if !ok { + return nil, fmt.Errorf("subagent: agent %q does not implement InvokableTool", name) + } + subAgentToolMap[name] = it + } + + toolName := config.ToolName + if toolName == "" { + toolName = agentToolName + } + + descGen := defaultAgentToolDescription[M] + if config.ToolDescriptionGenerator != nil { + descGen = config.ToolDescriptionGenerator + } + // The sub-agent set is fixed at construction, so the description is computed once. + desc, err := descGen(ctx, config.SubAgents) + if err != nil { + return nil, err + } + + // With a Manager, the tool exposes run_in_background and routes through the + // Manager; without one it is a plain foreground spawn. + var at tool.BaseTool + if config.Background != nil && config.Background.Manager != nil { + at, err = newManagedAgentTool(config.Background.Manager, subAgentToolMap, config.Background.OutputStore, config.Background.OutputDir, toolName, desc) + } else { + at, err = newAgentTool(subAgentToolMap, toolName, desc) + } + if err != nil { + return nil, err + } + + tools := []tool.BaseTool{at} + + // Build system prompt. + var instruction string + if config.SystemPrompt != nil { + instruction = *config.SystemPrompt + } else { + instruction = internal.SelectPrompt(internal.I18nPrompts{ + English: agentToolPrompt, + Chinese: agentToolPromptChinese, + }) + if config.Background != nil && config.Background.Manager != nil { + instruction += internal.SelectPrompt(internal.I18nPrompts{ + English: agentToolBackgroundPrompt, + Chinese: agentToolBackgroundPromptChinese, + }) + } + } + + return &typedSubagentMiddleware[M]{ + tools: tools, + instruction: instruction, + }, nil +} + +type typedSubagentMiddleware[M adk.MessageType] struct { + adk.TypedBaseChatModelAgentMiddleware[M] + tools []tool.BaseTool + instruction string +} + +// BeforeAgent injects sub-agent tools and instructions into the agent context. +func (m *typedSubagentMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { + if runCtx == nil { + return ctx, runCtx, nil + } + + nRunCtx := *runCtx + nRunCtx.Instruction += "\n" + m.instruction + nRunCtx.Tools = append(nRunCtx.Tools, m.tools...) + return ctx, &nRunCtx, nil +} + +func validate[M adk.MessageType](ctx context.Context, c *TypedConfig[M]) error { + if len(c.SubAgents) == 0 { + return fmt.Errorf("subagent: SubAgents must not be empty") + } + + names := make(map[string]struct{}, len(c.SubAgents)) + for _, a := range c.SubAgents { + name := a.Name(ctx) + if _, exists := names[name]; exists { + return fmt.Errorf("subagent: duplicate agent name %q", name) + } + names[name] = struct{}{} + } + + return nil +} diff --git a/adk/middlewares/subagent/middleware_test.go b/adk/middlewares/subagent/middleware_test.go new file mode 100644 index 000000000..0c5c11113 --- /dev/null +++ b/adk/middlewares/subagent/middleware_test.go @@ -0,0 +1,488 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package subagent + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// --- Mock Agent --- + +func intPtr(v int) *int { return &v } + +// anyRunning reports whether the manager still has a task in StatusRunning, +// derived from the public List() snapshot. +func anyRunning(m *backgroundtask.Manager) bool { + for _, t := range m.List() { + if t.Status == backgroundtask.StatusRunning { + return true + } + } + return false +} + +func waitAllTasks(t *testing.T, m *backgroundtask.Manager) { + t.Helper() + require.Eventually(t, func() bool { + return !anyRunning(m) + }, time.Second, 10*time.Millisecond) +} + +type mockAgent struct { + name string + desc string + // runFunc allows custom behavior in Run. + runFunc func(ctx context.Context, input *adk.AgentInput) string +} + +func (m *mockAgent) Name(_ context.Context) string { + return m.name +} + +func (m *mockAgent) Description(_ context.Context) string { + return m.desc +} + +func (m *mockAgent) Run(ctx context.Context, input *adk.AgentInput, options ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + result := m.desc // default: return description as result + if m.runFunc != nil { + result = m.runFunc(ctx, input) + } + + gen.Send(adk.EventFromMessage(schema.UserMessage(result), nil, schema.User, "")) + gen.Close() + return iter +} + +// --- Config Validation Tests --- + +func TestConfigValidation_EmptySubAgents(t *testing.T) { + _, err := New(context.Background(), &Config{ + SubAgents: nil, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "must not be empty") +} + +func TestConfigValidation_DuplicateNames(t *testing.T) { + _, err := New(context.Background(), &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "agent1", desc: "first"}, + &mockAgent{name: "agent1", desc: "second"}, + }, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "duplicate") +} + +// --- Middleware BeforeAgent Tests --- + +func TestBeforeAgent_InjectsToolsAndInstruction(t *testing.T) { + ctx := context.Background() + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "researcher", desc: "researches things"}, + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base instruction", + } + + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + // Instruction should be appended. + assert.Contains(t, newRunCtx.Instruction, "base instruction") + assert.Contains(t, newRunCtx.Instruction, "agent") + + // Agent tool should be injected. + assert.Len(t, newRunCtx.Tools, 1) +} + +func TestBeforeAgent_NilRunCtx(t *testing.T) { + ctx := context.Background() + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "helper", desc: "helps"}, + }, + }) + require.NoError(t, err) + + newCtx, newRunCtx, err := mw.BeforeAgent(ctx, nil) + require.NoError(t, err) + assert.Nil(t, newRunCtx) + assert.Equal(t, ctx, newCtx) +} + +func TestBeforeAgent_WithManager_InjectsAgentToolOnly(t *testing.T) { + ctx := context.Background() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + closeCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + _ = mgr.Close(closeCtx) + }() + + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "worker", desc: "does work"}, + }, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + } + + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + // Only the agent tool is injected here; task_output/task_stop are owned by + // the backgroundtask control middleware. + assert.Len(t, newRunCtx.Tools, 1) + + // Instruction should include the background-support prompt. + assert.Contains(t, newRunCtx.Instruction, "background") +} + +func TestBeforeAgent_CustomSystemPrompt(t *testing.T) { + ctx := context.Background() + customPrompt := "custom prompt" + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "helper", desc: "helps"}, + }, + SystemPrompt: &customPrompt, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{ + Instruction: "base", + } + + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + assert.Contains(t, newRunCtx.Instruction, "custom prompt") +} + +// --- Agent Tool Tests --- + +func TestAgentTool_ForegroundRouting(t *testing.T) { + ctx := context.Background() + a1 := &mockAgent{name: "agent1", desc: "desc of agent 1"} + a2 := &mockAgent{name: "agent2", desc: "desc of agent 2"} + + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{a1, a2}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + // Get the agent tool. + require.Len(t, newRunCtx.Tools, 1) + + // Use the tool directly. + at := newRunCtx.Tools[0].(tool.InvokableTool) + + result, err := at.InvokableRun(ctx, `{"subagent_type":"agent1","prompt":"test task","description":"test"}`) + require.NoError(t, err) + assert.Equal(t, "desc of agent 1", result) + + result, err = at.InvokableRun(ctx, `{"subagent_type":"agent2","prompt":"test task","description":"test"}`) + require.NoError(t, err) + assert.Equal(t, "desc of agent 2", result) +} + +func TestAgentTool_NotFound(t *testing.T) { + ctx := context.Background() + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "agent1", desc: "desc"}, + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + at := newRunCtx.Tools[0].(tool.InvokableTool) + _, err = at.InvokableRun(ctx, `{"subagent_type":"nonexistent","prompt":"test","description":"test"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "not found") +} + +func TestAgentTool_Background(t *testing.T) { + ctx := context.Background() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + closeCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(closeCtx) + }() + + slowAgent := &mockAgent{ + name: "slow", + desc: "slow agent", + runFunc: func(ctx context.Context, input *adk.AgentInput) string { + time.Sleep(50 * time.Millisecond) + return "slow result" + }, + } + + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{slowAgent}, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + at := newRunCtx.Tools[0].(tool.InvokableTool) + result, err := at.InvokableRun(ctx, `{"subagent_type":"slow","prompt":"bg task detail","description":"bg task","run_in_background":true}`) + require.NoError(t, err) + assert.Contains(t, result, "running in background") + assert.True(t, anyRunning(mgr)) + + // Wait for the background task to complete, then inspect final state. + waitAllTasks(t, mgr) + + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + assert.Equal(t, "slow result", tasks[0].Result) +} + +func TestAgentTool_Info(t *testing.T) { + ctx := context.Background() + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "helper", desc: "helps with tasks"}, + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + info, err := newRunCtx.Tools[0].Info(ctx) + require.NoError(t, err) + assert.Equal(t, agentToolName, info.Name) + assert.Contains(t, info.Desc, "helper") + assert.Contains(t, info.Desc, "helps with tasks") +} + +func TestAgentTool_CustomName(t *testing.T) { + ctx := context.Background() + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{ + &mockAgent{name: "helper", desc: "helps"}, + }, + ToolName: "task", + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + info, err := newRunCtx.Tools[0].Info(ctx) + require.NoError(t, err) + assert.Equal(t, "task", info.Name) +} + +// --- Foreground with Manager tracking --- + +func TestAgentTool_ForegroundWithTaskMgr(t *testing.T) { + ctx := context.Background() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + closeCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(closeCtx) + }() + + agent := &mockAgent{name: "fast", desc: "fast agent"} + + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{agent}, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + at := newRunCtx.Tools[0].(tool.InvokableTool) + + // Foreground run with TaskMgr: should block and return result. + result, err := at.InvokableRun(ctx, `{"subagent_type":"fast","prompt":"foreground task detail","description":"foreground task"}`) + require.NoError(t, err) + assert.Equal(t, "fast agent", result) + + // Task should be completed in TaskMgr. + assert.False(t, anyRunning(mgr)) + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + assert.Equal(t, "fast agent", tasks[0].Result) +} + +// With OutputStore and OutputDir configured, a completed managed agent run writes +// its final result to the task's output file. +func TestAgentTool_WritesOutputFile(t *testing.T) { + ctx := context.Background() + backend := filesystem.NewInMemoryBackend() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{}) + defer func() { + closeCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(closeCtx) + }() + + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{&mockAgent{name: "fast", desc: "fast agent"}}, + Background: &BackgroundConfig{ + Manager: mgr, + OutputStore: backend, + OutputDir: "/tasks", + }, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + at := newRunCtx.Tools[0].(tool.InvokableTool) + + _, err = at.InvokableRun(ctx, `{"subagent_type":"fast","prompt":"task detail","description":"task"}`) + require.NoError(t, err) + + tasks := mgr.List() + require.Len(t, tasks, 1) + path := tasks[0].OutputFile + require.NotEmpty(t, path) + + got, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: path}) + require.NoError(t, err) + assert.Equal(t, "fast agent", got.Content) +} + +// --- Auto-background --- + +func TestAgentTool_AutoBackground(t *testing.T) { + ctx := context.Background() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{ + ForegroundTimeoutMs: intPtr(50), // 50ms deadline + ShouldAutoBackground: func(context.Context, *backgroundtask.Task) bool { return true }, + }) + defer func() { + closeCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(closeCtx) + }() + + slowAgent := &mockAgent{ + name: "slow", + desc: "slow agent", + runFunc: func(ctx context.Context, input *adk.AgentInput) string { + time.Sleep(200 * time.Millisecond) + return "slow result" + }, + } + + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{slowAgent}, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + at := newRunCtx.Tools[0].(tool.InvokableTool) + + // Should auto-background after 50ms since agent takes 200ms. + result, err := at.InvokableRun(ctx, `{"subagent_type":"slow","prompt":"auto-bg task detail","description":"auto-bg task"}`) + require.NoError(t, err) + assert.Contains(t, result, "running in background") + + // Task should still be running. + assert.True(t, anyRunning(mgr)) + + // Wait for completion. + waitAllTasks(t, mgr) + + tasks := mgr.List() + require.Len(t, tasks, 1) + assert.Equal(t, backgroundtask.StatusCompleted, tasks[0].Status) + assert.Equal(t, "slow result", tasks[0].Result) +} + +func TestAgentTool_AutoBackground_FastAgent(t *testing.T) { + ctx := context.Background() + mgr := backgroundtask.New(context.Background(), &backgroundtask.Config{ForegroundTimeoutMs: intPtr(5000)}) // 5s timeout, agent finishes instantly + defer func() { + closeCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = mgr.Close(closeCtx) + }() + + fastAgent := &mockAgent{name: "fast", desc: "fast agent"} + + mw, err := New(ctx, &Config{ + SubAgents: []adk.Agent{fastAgent}, + Background: &BackgroundConfig{Manager: mgr}, + }) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext[*schema.Message]{} + _, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + require.NoError(t, err) + + at := newRunCtx.Tools[0].(tool.InvokableTool) + + // Fast agent completes before timeout — should return foreground result. + result, err := at.InvokableRun(ctx, `{"subagent_type":"fast","prompt":"fast task detail","description":"fast task"}`) + require.NoError(t, err) + assert.Equal(t, "fast agent", result) + assert.False(t, anyRunning(mgr)) +} diff --git a/adk/middlewares/subagent/prompt.go b/adk/middlewares/subagent/prompt.go new file mode 100644 index 000000000..befed3a35 --- /dev/null +++ b/adk/middlewares/subagent/prompt.go @@ -0,0 +1,148 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package subagent provides a ChatModelAgentMiddleware that injects Agent, TaskOutput, +// and TaskStop tools for spawning and managing sub-agents. +package subagent + +// This file contains prompt templates and tool descriptions for the subagent middleware. + +const ( + agentToolPrompt = ` +# Agent Tool + +You have access to an 'agent' tool to launch specialized agents that handle isolated tasks autonomously. Each agent invocation starts fresh — provide a complete task description. + +When to use the agent tool: +- When a task is complex and multi-step, and can be fully delegated in isolation +- When a task is independent of other tasks and can run in parallel +- When a task requires focused reasoning or heavy token/context usage that would bloat the orchestrator thread +- When you only care about the output of the subagent, and not the intermediate steps (e.g. performing research then returning a synthesized report) + +When NOT to use the agent tool: +- If you need to see the intermediate reasoning or steps (the agent tool hides them) +- If the task is trivial (a few tool calls or simple lookup) +- If delegating does not reduce token usage, complexity, or context switching + +## Usage Notes +- Whenever possible, parallelize the work. Launch multiple agents concurrently by issuing multiple tool calls within a single response. This saves time for the user. +- Always include a short description (3-5 words) summarizing what the agent will do. +- The agent's outputs should generally be trusted. +- Clearly tell the agent whether you expect it to write code or just to do research (search, file reads, etc.), since it is not aware of the user's intent. +- If the agent description mentions that it should be used proactively, then you should try your best to use it without the user having to ask for it first. +- If the user specifies that they want you to run agents "in parallel", you MUST issue multiple Agent tool calls within a single response. + +## Writing the prompt + +Brief the agent like a smart colleague who just walked into the room — it hasn't seen this conversation, doesn't know what you've tried, doesn't understand why this task matters. +- Explain what you're trying to accomplish and why. +- Describe what you've already learned or ruled out. +- Give enough context about the surrounding problem that the agent can make judgment calls rather than just following a narrow instruction. +- If you need a short response, say so ("report in under 200 words"). +- Lookups: hand over the exact command. Investigations: hand over the question — prescribed steps become dead weight when the premise is wrong. + +Terse command-style prompts produce shallow, generic work. + +**Never delegate understanding.** Don't write "based on your findings, fix the bug" or "based on the research, implement it." Those phrases push synthesis onto the agent instead of doing it yourself. Write prompts that prove you understood: include file paths, line numbers, what specifically to change. +` + + agentToolPromptChinese = ` +# Agent 工具 + +你可以使用 'agent' 工具启动专门的智能体来自主处理独立任务。每次智能体调用都从零开始——请提供完整的任务描述。 + +何时使用 agent 工具： +- 当任务复杂且包含多个步骤，并且可以完全独立委托时 +- 当任务独立于其他任务并且可以并行运行时 +- 当任务需要集中推理或大量 token/上下文使用，这会使编排器线程膨胀时 +- 当你只关心子智能体的输出，而不关心中间步骤时（例如执行大量研究然后返回综合报告） + +何时不使用 agent 工具： +- 如果你需要查看中间推理或步骤（agent 工具会隐藏它们） +- 如果任务很简单（几个工具调用或简单查找） +- 如果委托不会减少 token 使用、复杂性或上下文切换 + +## 使用注意事项 +- 尽可能并行化工作。通过在一条消息中使用多个工具调用来同时启动多个智能体。这为用户节省了时间。 +- 始终包含一个简短的描述（3-5 个词）来概括智能体要做的事情。 +- 智能体的输出通常应该被信任。 +- 明确告诉智能体你期望它编写代码还是只是进行研究（搜索、文件读取等），因为它不知道用户的意图。 +- 如果智能体描述提到应该主动使用它，那么你应该尽力主动使用它。 +- 如果用户指定他们希望你"并行"运行智能体，你必须在一次回复中发起多个 Agent 工具调用。 + +## 编写提示词 + +像给一个刚走进房间的聪明同事做简报一样对待智能体——它没有看过这段对话，不知道你尝试过什么，不了解为什么这个任务重要。 +- 解释你要完成什么以及为什么。 +- 描述你已经了解到或排除的内容。 +- 提供足够的背景上下文，使智能体能够做出判断而不只是执行狭隘的指令。 +- 如果你需要简短的回复，请说明（"200 字以内报告"）。 +- 查找任务：给出确切的命令。调查任务：给出问题——预设步骤在前提错误时会成为负担。 + +简短的命令式提示词会产生浅层、泛化的结果。 + +**不要把"理解问题"这一步交给子智能体。**不要写"根据你的发现修复这个 bug"或"根据研究来实现它"——这类写法把本该由你完成的分析与综合推给了子智能体。要写出能证明你已经理解的提示词：包含文件路径、行号、具体要改什么。 +` + + agentToolDescription = `Launch a new agent to handle complex, multi-step tasks. Each agent type has specific capabilities and tools available to it. + +When using the agent tool, specify a subagent_type parameter to select which agent type to use. + +Available agent types and the tools they have access to: +{other_agents} + +## When to use + +Reach for this when the task matches an available agent type, when you have independent work to run in parallel, or when answering would mean reading across several files — delegate it and you keep the conclusion, not the file dumps. For a single-fact lookup where you already know the file, symbol, or value, search directly. Once you've delegated a search, don't also run it yourself — wait for the result. + +- The agent's final message is returned to you as the tool result; it is not shown to the user — relay what matters. +- Each agent call starts fresh, so give a complete, self-contained task description. +` + + agentToolDescriptionChinese = `启动新智能体来处理复杂的多步骤任务。每种智能体类型都有特定的能力和可用的工具。 + +使用 agent 工具时，指定 subagent_type 参数来选择要使用的智能体类型。 + +可用的智能体类型及其可访问的工具： +{other_agents} + +## 何时使用 + +当任务匹配某个可用的智能体类型、当你有可以并行处理的独立工作、或者当回答问题需要跨多个文件阅读时——把它委托出去，你只需保留结论，而无需处理大量文件内容。对于你已经知道文件、符号或具体值的单点查找，直接自己搜索即可。一旦你把某个搜索委托出去，就不要自己再重复执行——等待它的结果。 + +- 智能体的最终消息会作为工具结果返回给你；它不会展示给用户——请转述其中重要的内容。 +- 每次智能体调用都是全新开始，因此请提供完整、自包含的任务描述。 +` + + agentToolBackgroundPrompt = ` +## Running agents in the background +- Set run_in_background=true to run an agent in the background. It keeps running after the tool + call returns, and you will be notified when it completes. Do not block waiting on it — continue + with other work, and use the task_output tool to check its status or retrieve its result by + task_id when you need it. +- Use foreground (the default) when you need the agent's result before you can proceed; use + background when you have genuinely independent work to do in parallel. +- Use the task_stop tool to cancel a background agent by task_id. +` + + agentToolBackgroundPromptChinese = ` +## 在后台运行智能体 +- 设置 run_in_background=true 可在后台运行智能体。它在工具调用返回后会继续运行，完成时你将收到通知。 + 不要为等待它而阻塞——请继续处理其他工作，并在需要时使用 task_output 工具通过 task_id 查询其状态或获取结果。 +- 当你需要智能体的结果才能继续时使用前台（默认）；当你有真正独立的工作可以并行完成时使用后台。 +- 使用 task_stop 工具通过 task_id 取消后台智能体。 +` +) diff --git a/adk/middlewares/summarization/finalizer_builder.go b/adk/middlewares/summarization/finalizer_builder.go index 0d10b7682..7e79b2b8f 100644 --- a/adk/middlewares/summarization/finalizer_builder.go +++ b/adk/middlewares/summarization/finalizer_builder.go @@ -263,7 +263,7 @@ func extractSkillInfos[M adk.MessageType](messages []M, skillTool string) ([]*sk } skills = append(skills, &skillInfo{ Name: arg.Skill, - Content: m.Content, + Content: messageUserTextContent(m), }) } diff --git a/adk/middlewares/summarization/prompt.go b/adk/middlewares/summarization/prompt.go index 086017e90..13be8f814 100644 --- a/adk/middlewares/summarization/prompt.go +++ b/adk/middlewares/summarization/prompt.go @@ -22,7 +22,7 @@ import ( "github.com/cloudwego/eino/adk/internal" ) -var allUserMessagesTagRegex = regexp.MustCompile(`(?s).*`) +var allUserMessagesTagRegex = regexp.MustCompile(`(?s).*?`) func getSystemInstruction() string { return internal.SelectPrompt(internal.I18nPrompts{ diff --git a/adk/middlewares/summarization/summarization.go b/adk/middlewares/summarization/summarization.go index a99bf528f..189e74822 100644 --- a/adk/middlewares/summarization/summarization.go +++ b/adk/middlewares/summarization/summarization.go @@ -354,6 +354,24 @@ func (m *TypedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state afterState := *state afterState.Messages = finalMsgs + // Emit a session mutation event so the persisted event log reflects the new + // message state at the summarization boundary. Independent of EmitInternalEvents. + // Error is ignored: when not in an execution context (e.g. unit tests), the + // event simply has no consumer. + // + // The emitted durable payload strips runtime-generated leading system messages + // because they are recalculated on each run via applyBeforeAgent -> GenModelInput + // and should not be reconstructed from session history. + msgs := stripRuntimeGeneratedLeadingSystemMessages(afterState.Messages) + _ = adk.TypedSendEvent(ctx, &adk.TypedAgentEvent[M]{ + SessionEventVariant: &adk.SessionEventVariant[M]{ + Event: &adk.SessionEvent[M]{ + Kind: adk.SessionEventMessagesReplaced, + MessagesReplaced: &msgs, + }, + }, + }) + return ctx, &afterState, nil } @@ -961,6 +979,8 @@ func (c *TriggerCondition) check() error { // Generic helper functions // ============================================================================ +const extraKeyRuntimeGeneratedSystemMessage = "_eino_adk_runtime_generated_system_message" + func isSystemRole[M adk.MessageType](msg M) bool { switch m := any(msg).(type) { case *schema.Message: @@ -971,6 +991,32 @@ func isSystemRole[M adk.MessageType](msg M) bool { panic("unreachable") } +func isMarkedRuntimeGeneratedSystemMessage[M adk.MessageType](msg M) bool { + if !isSystemRole(msg) { + return false + } + extra := getMsgExtra(msg) + if extra == nil { + return false + } + v, ok := extra[extraKeyRuntimeGeneratedSystemMessage] + if !ok { + return false + } + b, ok := v.(bool) + return ok && b +} + +func stripRuntimeGeneratedLeadingSystemMessages[M adk.MessageType](msgs []M) []M { + i := 0 + for i < len(msgs) && isMarkedRuntimeGeneratedSystemMessage(msgs[i]) { + i++ + } + out := make([]M, 0, len(msgs)-i) + out = append(out, msgs[i:]...) + return out +} + func isUserRole[M adk.MessageType](msg M) bool { switch m := any(msg).(type) { case *schema.Message: @@ -981,22 +1027,26 @@ func isUserRole[M adk.MessageType](msg M) bool { panic("unreachable") } +func messageUserTextContent(m *schema.Message) string { + if m == nil { + return "" + } + var parts []string + for _, part := range m.UserInputMultiContent { + if part.Type == schema.ChatMessagePartTypeText && part.Text != "" { + parts = append(parts, part.Text) + } + } + if len(parts) > 0 { + return strings.Join(parts, "\n") + } + return m.Content +} + func getUserMsgTextContent[M adk.MessageType](msg M) string { switch m := any(msg).(type) { case *schema.Message: - if m == nil { - return "" - } - var parts []string - for _, part := range m.UserInputMultiContent { - if part.Type == schema.ChatMessagePartTypeText && part.Text != "" { - parts = append(parts, part.Text) - } - } - if len(parts) > 0 { - return strings.Join(parts, "\n") - } - return m.Content + return messageUserTextContent(m) case *schema.AgenticMessage: if m == nil { diff --git a/adk/middlewares/summarization/summarization_attack_review_test.go b/adk/middlewares/summarization/summarization_attack_review_test.go new file mode 100644 index 000000000..5f0bffe67 --- /dev/null +++ b/adk/middlewares/summarization/summarization_attack_review_test.go @@ -0,0 +1,573 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package summarization + +import ( + "context" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +// ============================================================================= +// Attack tests for getAssistantTextContent +// ============================================================================= + +func TestAttack_GetAssistantTextContent_BothContentAndMultiContent(t *testing.T) { + // When both Content and AssistantGenMultiContent are populated, + // the function should prefer AssistantGenMultiContent. + msg := &schema.Message{ + Role: schema.Assistant, + Content: "plain content fallback", + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "multi part 1"}, + {Type: schema.ChatMessagePartTypeText, Text: "multi part 2"}, + }, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "multi part 1\nmulti part 2", result) + assert.NotContains(t, result, "plain content fallback", + "should prefer AssistantGenMultiContent over Content field") +} + +func TestAttack_GetAssistantTextContent_FallbackToContent(t *testing.T) { + // When AssistantGenMultiContent is empty, should fall back to Content. + msg := &schema.Message{ + Role: schema.Assistant, + Content: "fallback content", + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "fallback content", result) +} + +func TestAttack_GetAssistantTextContent_EmptyMultiContentParts(t *testing.T) { + // When AssistantGenMultiContent has parts but all have empty Text, + // the function should fall back to Content. + msg := &schema.Message{ + Role: schema.Assistant, + Content: "should use this", + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: ""}, + {Type: schema.ChatMessagePartTypeImageURL}, // non-text type + }, + } + + result := getAssistantTextContent(msg) + // Empty text parts are filtered, so no parts collected → falls back to Content + assert.Equal(t, "should use this", result) +} + +func TestAttack_GetAssistantTextContent_MultiContentWithNonTextTypes(t *testing.T) { + // Non-text parts in AssistantGenMultiContent should be ignored. + msg := &schema.Message{ + Role: schema.Assistant, + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeImageURL}, + {Type: schema.ChatMessagePartTypeText, Text: "actual text"}, + {Type: schema.ChatMessagePartTypeReasoning, Reasoning: &schema.MessageOutputReasoning{Text: "reasoning"}}, + }, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "actual text", result, "should only extract text parts") +} + +func TestAttack_GetAssistantTextContent_AgenticMessage_NilBlocks(t *testing.T) { + // AgenticMessage with nil blocks in ContentBlocks should not panic. + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + nil, + schema.NewContentBlock(&schema.AssistantGenText{Text: "hello"}), + nil, + schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}), + }, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "hello\nworld", result) +} + +func TestAttack_GetAssistantTextContent_AgenticMessage_NonTextBlocks(t *testing.T) { + // AgenticMessage with non-text blocks (tool calls, images, etc.) should only get text. + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolCall{Name: "tool1", Arguments: "{}"}), + schema.NewContentBlock(&schema.AssistantGenText{Text: "response text"}), + schema.NewContentBlock(&schema.Reasoning{Text: "reasoning text"}), + }, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "response text", result, "should only extract AssistantGenText blocks") +} + +func TestAttack_GetAssistantTextContent_AgenticMessage_EmptyBlocks(t *testing.T) { + // AgenticMessage with empty ContentBlocks should return empty string. + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{}, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "", result) +} + +func TestAttack_GetAssistantTextContent_AgenticMessage_NilAssistantGenText(t *testing.T) { + // Block with Type == AssistantGenText but nil AssistantGenText field. + // The code checks `block.AssistantGenText != nil` so this should be safe. + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + {Type: schema.ContentBlockTypeAssistantGenText, AssistantGenText: nil}, + schema.NewContentBlock(&schema.AssistantGenText{Text: "valid"}), + }, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "valid", result) +} + +// ============================================================================= +// Attack tests for postProcessSummary with edge-case contextMsgs +// ============================================================================= + +func TestAttack_PostProcessSummary_EmptyContextMsgs(t *testing.T) { + // When contextMsgs is empty (len==0), replaceUserMessagesInSummary is skipped. + ctx := context.Background() + + summaryContent := "Summary with old content tag" + result, err := postProcessSummary(ctx, &postProcessSummaryParams[*schema.Message]{ + contextMsgs: nil, + summaryContent: summaryContent, + }) + require.NoError(t, err) + + // The tag should NOT be replaced because contextMsgs is empty + text := getUserMsgTextContent(result) + assert.Contains(t, text, "old content") +} + +func TestAttack_PostProcessSummary_AllContextMsgsAreSummaries(t *testing.T) { + // contextMsgs is non-empty but all messages have contentTypeSummary. + // replaceUserMessagesInSummary WILL be called (len > 0), but inside it + // all messages are filtered out because they have summary content type. + // The function should gracefully return the original summary text unchanged. + ctx := context.Background() + + summaryMsg := &schema.Message{ + Role: schema.User, + Content: "previous summary content", + Extra: map[string]any{extraKeyContentType: string(contentTypeSummary)}, + } + + summaryContent := "New summary with placeholder" + result, err := postProcessSummary(ctx, &postProcessSummaryParams[*schema.Message]{ + contextMsgs: []*schema.Message{summaryMsg}, + summaryContent: summaryContent, + }) + require.NoError(t, err) + + // Since all msgs are summaries, hasUserMsgs is false, so original text is preserved. + text := getUserMsgTextContent(result) + assert.Contains(t, text, "placeholder", + "tag should not be replaced when all context msgs are summaries") +} + +func TestAttack_PostProcessSummary_ContextMsgsNoUserMessages(t *testing.T) { + // contextMsgs has messages but none are user role. + ctx := context.Background() + + assistantMsg := &schema.Message{ + Role: schema.Assistant, + Content: "assistant response", + } + + summaryContent := "Summary content" + result, err := postProcessSummary(ctx, &postProcessSummaryParams[*schema.Message]{ + contextMsgs: []*schema.Message{assistantMsg}, + summaryContent: summaryContent, + }) + require.NoError(t, err) + + text := getUserMsgTextContent(result) + // No user messages found, so original tag preserved + assert.Contains(t, text, "content") +} + +// ============================================================================= +// Attack tests for buildInternalFinalizer + DefaultFinalize parity +// ============================================================================= + +func TestAttack_BuildInternalFinalizer_DefaultFinalize_Parity(t *testing.T) { + // When TranscriptFilePath is empty, buildInternalFinalizer and DefaultFinalize + // should produce identical results. + ctx := context.Background() + + systemMsg := schema.SystemMessage("You are a helpful assistant.") + userMsg := &schema.Message{Role: schema.User, Content: "Hello, please help me."} + assistantReply := &schema.Message{ + Role: schema.Assistant, + Content: "summary of conversation", + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "summary of conversation"}, + }, + } + + originalMsgs := []*schema.Message{systemMsg, userMsg} + + cfg := &TypedConfig[*schema.Message]{ + TranscriptFilePath: "", + } + + internalFinalizer := buildInternalFinalizer(cfg) + + result1, err := internalFinalizer(ctx, originalMsgs, assistantReply) + require.NoError(t, err) + + result2, err := DefaultFinalize(ctx, originalMsgs, assistantReply) + require.NoError(t, err) + + require.Equal(t, len(result1), len(result2), "should produce same number of messages") + for i := range result1 { + text1 := getUserMsgTextContent(result1[i]) + text2 := getUserMsgTextContent(result2[i]) + assert.Equal(t, text1, text2, "message %d content should be identical", i) + } +} + +func TestAttack_BuildInternalFinalizer_WithTranscriptPath(t *testing.T) { + // With TranscriptFilePath set, buildInternalFinalizer should include transcript path + // instruction, while DefaultFinalize should NOT include it. + ctx := context.Background() + + userMsg := &schema.Message{Role: schema.User, Content: "hello"} + assistantReply := &schema.Message{ + Role: schema.Assistant, + Content: "summary text", + } + originalMsgs := []*schema.Message{userMsg} + + cfg := &TypedConfig[*schema.Message]{ + TranscriptFilePath: "/path/to/transcript.md", + } + + internalFinalizer := buildInternalFinalizer(cfg) + result1, err := internalFinalizer(ctx, originalMsgs, assistantReply) + require.NoError(t, err) + + result2, err := DefaultFinalize(ctx, originalMsgs, assistantReply) + require.NoError(t, err) + + text1 := getUserMsgTextContent(result1[0]) + text2 := getUserMsgTextContent(result2[0]) + + assert.Contains(t, text1, "/path/to/transcript.md", + "internal finalizer should include transcript path") + assert.NotContains(t, text2, "/path/to/transcript.md", + "DefaultFinalize should NOT include transcript path") +} + +// ============================================================================= +// Attack tests for token budget overflow +// ============================================================================= + +func TestAttack_TokenBudgetOverflow_SingleLargeMessage(t *testing.T) { + // A single user message with >30000 tokens (>120000 chars at 4 chars/token). + // The trimming logic should handle this via defaultTypedTrimUserMessage. + ctx := context.Background() + + // Create a message much larger than 30000 tokens (> 120000 chars) + largeContent := strings.Repeat("x", 150000) // ~37500 tokens + + userMsg := &schema.Message{Role: schema.User, Content: largeContent} + summaryText := "Summary placeholder" + + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: []*schema.Message{userMsg}, + summaryText: summaryText, + }) + require.NoError(t, err) + + // Since there's only 1 user message, selected = userMsgs (no trimming in that branch) + // The code takes len(userMsgs)==1 as a special case: selected = userMsgs directly. + assert.Contains(t, result, "") + assert.Contains(t, result, "") +} + +func TestAttack_TokenBudgetOverflow_MultipleMessagesExceedBudget(t *testing.T) { + // Multiple user messages where each exceeds 30000 tokens. + // The trimming should kick in for the second message that crosses the budget. + ctx := context.Background() + + // Each message ~10000 tokens (40000 chars); 4 of them = 40000 tokens > 30000 budget + msgContent := strings.Repeat("a", 40000) + msgs := make([]*schema.Message, 4) + for i := range msgs { + msgs[i] = &schema.Message{Role: schema.User, Content: msgContent} + } + + summaryText := "Summary old" + + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: msgs, + summaryText: summaryText, + }) + require.NoError(t, err) + + // The result should contain the replacement and a note about cleared messages + assert.Contains(t, result, "") + assert.Contains(t, result, "") +} + +func TestAttack_TokenBudgetOverflow_TrimUserMessage(t *testing.T) { + // Verify defaultTypedTrimUserMessage with remaining budget > 0 produces truncated content. + largeContent := strings.Repeat("y", 200000) // ~50000 tokens + msg := &schema.Message{Role: schema.User, Content: largeContent} + + trimmed := defaultTypedTrimUserMessage(msg, 100) // very small remaining budget + text := getUserMsgTextContent(trimmed) + assert.NotEmpty(t, text, "trimmed message should not be empty") + assert.Less(t, len(text), len(largeContent), "trimmed should be shorter") +} + +func TestAttack_TokenBudgetOverflow_TrimUserMessageZeroBudget(t *testing.T) { + // With 0 remaining tokens, defaultTypedTrimUserMessage should return zero. + msg := &schema.Message{Role: schema.User, Content: "hello world"} + + trimmed := defaultTypedTrimUserMessage[*schema.Message](msg, 0) + assert.Nil(t, trimmed, "zero budget should return nil message") +} + +// ============================================================================= +// Attack tests for newTypedSummaryMessage metadata +// ============================================================================= + +func TestAttack_NewTypedSummaryMessage_ExtraMetadata(t *testing.T) { + // Verify the summary message has the correct extraKeyContentType set + // so recursive summarization doesn't re-process it. + msg := newTypedSummaryMessage[*schema.Message]("test summary content") + + assert.NotNil(t, msg.Extra) + ct, ok := msg.Extra[extraKeyContentType].(string) + require.True(t, ok, "extra should contain content type key") + assert.Equal(t, string(contentTypeSummary), ct) +} + +func TestAttack_NewTypedSummaryMessage_AgenticExtraMetadata(t *testing.T) { + // Verify AgenticMessage variant also gets proper metadata. + msg := newTypedSummaryMessage[*schema.AgenticMessage]("test agentic summary") + + assert.NotNil(t, msg.Extra) + ct, ok := msg.Extra[extraKeyContentType].(string) + require.True(t, ok, "extra should contain content type key") + assert.Equal(t, string(contentTypeSummary), ct) +} + +func TestAttack_NewTypedSummaryMessage_IsFilteredBySummarizationCheck(t *testing.T) { + // Verify that typedGetContentType correctly identifies summary messages, + // ensuring they are skipped in replaceUserMessagesInSummary. + msg := newTypedSummaryMessage[*schema.Message]("summary content") + + ct := typedGetContentType(msg) + assert.Equal(t, contentTypeSummary, ct) +} + +// ============================================================================= +// Attack tests for appendSection concatenation correctness +// ============================================================================= + +func TestAttack_AppendSection_BothNonEmpty(t *testing.T) { + result := appendSection("first part", "second part") + assert.Equal(t, "first part\n\nsecond part", result) +} + +func TestAttack_AppendSection_BaseEmpty(t *testing.T) { + result := appendSection("", "only section") + assert.Equal(t, "only section", result) +} + +func TestAttack_AppendSection_SectionEmpty(t *testing.T) { + result := appendSection("only base", "") + assert.Equal(t, "only base", result) +} + +func TestAttack_AppendSection_BothEmpty(t *testing.T) { + result := appendSection("", "") + assert.Equal(t, "", result) +} + +func TestAttack_AppendSection_FinalMessageWellFormed(t *testing.T) { + // Simulate the actual postProcessSummary concatenation flow: + // preamble + content + continueInstruction + preamble := getSummaryPreamble() + content := "Summary body text" + continueInstr := getContinueInstruction() + + step1 := appendSection(preamble, content) + final := appendSection(step1, continueInstr) + + // Verify structure: preamble, double newline, content, double newline, continue + parts := strings.Split(final, "\n\n") + assert.GreaterOrEqual(t, len(parts), 3, + "final message should have at least 3 sections separated by double newlines") + assert.Equal(t, preamble, parts[0]) +} + +// ============================================================================= +// Attack tests for AgenticMessage path in getAssistantTextContent +// ============================================================================= + +func TestAttack_GetAssistantTextContent_AgenticMessage_AllNilBlocks(t *testing.T) { + // All blocks are nil — should not panic and return empty string. + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + nil, nil, nil, + }, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "", result) +} + +func TestAttack_GetAssistantTextContent_AgenticMessage_MixedBlocksWithEmptyText(t *testing.T) { + // Mix of valid and empty-text AssistantGenText blocks. + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: ""}), + schema.NewContentBlock(&schema.AssistantGenText{Text: "non-empty"}), + schema.NewContentBlock(&schema.AssistantGenText{Text: ""}), + schema.NewContentBlock(&schema.AssistantGenText{Text: "also valid"}), + }, + } + + result := getAssistantTextContent(msg) + // The code does NOT filter empty text for AgenticMessage — it joins all AssistantGenText.Text + // including empty ones with "\n" + assert.Contains(t, result, "non-empty") + assert.Contains(t, result, "also valid") +} + +func TestAttack_GetAssistantTextContent_AgenticMessage_OnlyToolCalls(t *testing.T) { + // Only tool call blocks, no text at all. + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolCall{Name: "read", Arguments: `{"path":"test"}`}), + schema.NewContentBlock(&schema.FunctionToolCall{Name: "write", Arguments: `{"content":"x"}`}), + }, + } + + result := getAssistantTextContent(msg) + assert.Equal(t, "", result, "should return empty when only tool calls present") +} + +// ============================================================================= +// Attack tests for DefaultFinalize end-to-end behavior +// ============================================================================= + +func TestAttack_DefaultFinalize_PreservesSystemMessages(t *testing.T) { + ctx := context.Background() + + sys1 := schema.SystemMessage("system prompt 1") + sys2 := schema.SystemMessage("system prompt 2") + userMsg := &schema.Message{Role: schema.User, Content: "user question"} + originalMsgs := []*schema.Message{sys1, sys2, userMsg} + + summary := &schema.Message{ + Role: schema.Assistant, + Content: "conversation summary", + } + + result, err := DefaultFinalize(ctx, originalMsgs, summary) + require.NoError(t, err) + + // First two should be system messages + require.GreaterOrEqual(t, len(result), 3) + assert.Equal(t, schema.System, result[0].Role) + assert.Equal(t, schema.System, result[1].Role) + // Last one should be the processed summary (user role with summary content type) + lastMsg := result[len(result)-1] + assert.Equal(t, schema.User, lastMsg.Role) + ct := typedGetContentType(lastMsg) + assert.Equal(t, contentTypeSummary, ct, "final message should be marked as summary") +} + +func TestAttack_DefaultFinalize_EmptySummaryContent(t *testing.T) { + // What happens if the model returned an empty summary? + ctx := context.Background() + + userMsg := &schema.Message{Role: schema.User, Content: "test"} + originalMsgs := []*schema.Message{userMsg} + + summary := &schema.Message{ + Role: schema.Assistant, + Content: "", + } + + _, err := DefaultFinalize(ctx, originalMsgs, summary) + require.Error(t, err, "empty summary content should return an error") + assert.Contains(t, err.Error(), "summary content is empty") +} + +// ============================================================================= +// Attack test for replaceUserMessagesInSummary with no tag +// ============================================================================= + +func TestAttack_ReplaceUserMessages_NoTag(t *testing.T) { + // If the summary doesn't contain the tag, + // the function should return the original text unchanged. + ctx := context.Background() + + userMsg := &schema.Message{Role: schema.User, Content: "hello"} + summaryText := "This is a summary without any tag markers." + + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: []*schema.Message{userMsg}, + summaryText: summaryText, + }) + require.NoError(t, err) + assert.Equal(t, summaryText, result) +} + +func TestAttack_ReplaceUserMessages_MultipleTagInstances(t *testing.T) { + // If there are multiple tags, only the LAST one should be replaced. + ctx := context.Background() + + userMsg := &schema.Message{Role: schema.User, Content: "my message"} + summaryText := "first middle second" + + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: []*schema.Message{userMsg}, + summaryText: summaryText, + }) + require.NoError(t, err) + + // First tag should be preserved, last one replaced + assert.Contains(t, result, "first", + "first tag should remain unchanged") + assert.Contains(t, result, "my message", "user message should appear in replacement") +} diff --git a/adk/middlewares/summarization/summarization_test.go b/adk/middlewares/summarization/summarization_test.go index d70f396b0..7771295f0 100644 --- a/adk/middlewares/summarization/summarization_test.go +++ b/adk/middlewares/summarization/summarization_test.go @@ -168,6 +168,90 @@ func TestMiddlewareBeforeModelRewriteState(t *testing.T) { assert.Equal(t, schema.User, newState.Messages[1].Role) }) + t.Run("marked runtime system messages stripped from MessagesReplaced event payload", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + assert.Equal(t, schema.System, msgs[0].Role) + return &schema.Message{ + Role: schema.Assistant, + Content: "Summary content", + }, nil + }).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + runtimeSys := schema.SystemMessage("runtime generated system") + setMsgExtra(runtimeSys, extraKeyRuntimeGeneratedSystemMessage, true) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + runtimeSys, + schema.UserMessage(strings.Repeat("a", 100)), + schema.AssistantMessage(strings.Repeat("b", 100), nil), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 2) + assert.Equal(t, schema.System, newState.Messages[0].Role) + assert.Equal(t, "runtime generated system", newState.Messages[0].Content) + + eventPayload := stripRuntimeGeneratedLeadingSystemMessages(newState.Messages) + assert.Len(t, eventPayload, 1) + assert.Equal(t, schema.User, eventPayload[0].Role) + }) + + t.Run("unmarked caller system messages preserved in both runtime and event payload", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + assert.Equal(t, schema.System, msgs[0].Role) + return &schema.Message{ + Role: schema.Assistant, + Content: "Summary content", + }, nil + }).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + callerSys := schema.SystemMessage("caller supplied system") + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + callerSys, + schema.UserMessage(strings.Repeat("a", 100)), + schema.AssistantMessage(strings.Repeat("b", 100), nil), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 2) + assert.Equal(t, schema.System, newState.Messages[0].Role) + assert.Equal(t, "caller supplied system", newState.Messages[0].Content) + + eventPayload := stripRuntimeGeneratedLeadingSystemMessages(newState.Messages) + assert.Len(t, eventPayload, 2) + assert.Equal(t, schema.System, eventPayload[0].Role) + assert.Equal(t, "caller supplied system", eventPayload[0].Content) + }) + t.Run("preserves multiple system messages", func(t *testing.T) { ctrl := gomock.NewController(t) cm := mockModel.NewMockBaseChatModel(ctrl) @@ -1425,11 +1509,10 @@ func TestPostProcessSummary(t *testing.T) { func TestEventHelpers(t *testing.T) { ctx := context.Background() - t.Run("emitEvent returns wrapped error outside execution context", func(t *testing.T) { + t.Run("emitEvent is no-op outside execution context", func(t *testing.T) { mw := &TypedMiddleware[*schema.Message]{cfg: &Config{}} err := mw.emitEvent(ctx, &CustomizedAction{Type: ActionTypeBeforeSummarize}) - assert.Error(t, err) - assert.Contains(t, err.Error(), "failed to send internal event") + assert.NoError(t, err) }) t.Run("emitGenerateSummaryEvent is skipped when internal events are disabled", func(t *testing.T) { @@ -1438,11 +1521,10 @@ func TestEventHelpers(t *testing.T) { assert.NoError(t, err) }) - t.Run("emitGenerateSummaryEvent returns wrapped error when enabled outside execution context", func(t *testing.T) { + t.Run("emitGenerateSummaryEvent is no-op when enabled outside execution context", func(t *testing.T) { mw := &TypedMiddleware[*schema.Message]{cfg: &Config{EmitInternalEvents: true}} err := mw.emitGenerateSummaryEvent(ctx, 1, GenerateSummaryPhasePrimary, schema.AssistantMessage("ok", nil), nil) - assert.Error(t, err) - assert.Contains(t, err.Error(), "failed to send internal event") + assert.NoError(t, err) }) } @@ -1937,7 +2019,7 @@ func TestSummarizationGeneric(t *testing.T) { }) } -func TestEmitInternalEvents_AgenticMessage_RequiresExecContext(t *testing.T) { +func TestEmitInternalEvents_AgenticMessage_NoopOutsideExecContext(t *testing.T) { ctx := context.Background() longContent := strings.Repeat("x", 800000) @@ -1967,9 +2049,12 @@ func TestEmitInternalEvents_AgenticMessage_RequiresExecContext(t *testing.T) { require.NoError(t, err) state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: msgs} - _, _, err = mw.BeforeModelRewriteState(ctx, state, nil) - assert.Error(t, err, "should error without exec context when EmitInternalEvents is true") - assert.Contains(t, err.Error(), "send internal event") + _, gotState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + require.NotNil(t, gotState) + require.Len(t, gotState.Messages, 2) + assert.Equal(t, schema.AgenticRoleTypeSystem, gotState.Messages[0].Role) + assert.Equal(t, schema.AgenticRoleTypeUser, gotState.Messages[1].Role) } func testSummarizationHelpers[M adk.MessageType](t *testing.T) { @@ -2243,3 +2328,162 @@ func TestGetAssistantTextContent(t *testing.T) { assert.Equal(t, "", got) }) } + +func TestStripRuntimeGeneratedLeadingSystemMessages(t *testing.T) { + t.Run("no system messages", func(t *testing.T) { + msgs := []*schema.Message{ + schema.UserMessage("hello"), + schema.AssistantMessage("hi", nil), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 2) + assert.Equal(t, "hello", result[0].Content) + }) + + t.Run("unmarked system message preserved", func(t *testing.T) { + sys := schema.SystemMessage("durable system") + msgs := []*schema.Message{ + sys, + schema.UserMessage("hello"), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 2) + assert.Equal(t, schema.System, result[0].Role) + assert.Equal(t, "durable system", result[0].Content) + }) + + t.Run("single marked system message stripped", func(t *testing.T) { + sys := schema.SystemMessage("runtime system") + setMsgExtra(sys, extraKeyRuntimeGeneratedSystemMessage, true) + msgs := []*schema.Message{ + sys, + schema.UserMessage("hello"), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 1) + assert.Equal(t, schema.User, result[0].Role) + assert.Equal(t, "hello", result[0].Content) + }) + + t.Run("multiple marked system messages stripped", func(t *testing.T) { + sys1 := schema.SystemMessage("runtime 1") + setMsgExtra(sys1, extraKeyRuntimeGeneratedSystemMessage, true) + sys2 := schema.SystemMessage("runtime 2") + setMsgExtra(sys2, extraKeyRuntimeGeneratedSystemMessage, true) + msgs := []*schema.Message{ + sys1, + sys2, + schema.UserMessage("hello"), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 1) + assert.Equal(t, schema.User, result[0].Role) + }) + + t.Run("mixed marked and unmarked - only prefix stripped", func(t *testing.T) { + marked := schema.SystemMessage("runtime") + setMsgExtra(marked, extraKeyRuntimeGeneratedSystemMessage, true) + unmarked := schema.SystemMessage("durable") + msgs := []*schema.Message{ + marked, + unmarked, + schema.UserMessage("hello"), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 2) + assert.Equal(t, schema.System, result[0].Role) + assert.Equal(t, "durable", result[0].Content) + }) + + t.Run("non-leading system message preserved", func(t *testing.T) { + sys := schema.SystemMessage("mid-system") + setMsgExtra(sys, extraKeyRuntimeGeneratedSystemMessage, true) + msgs := []*schema.Message{ + schema.UserMessage("first"), + sys, + schema.AssistantMessage("resp", nil), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 3) + assert.Equal(t, schema.User, result[0].Role) + assert.Equal(t, schema.System, result[1].Role) + }) + + t.Run("only marked system messages - empty result", func(t *testing.T) { + sys1 := schema.SystemMessage("runtime 1") + setMsgExtra(sys1, extraKeyRuntimeGeneratedSystemMessage, true) + sys2 := schema.SystemMessage("runtime 2") + setMsgExtra(sys2, extraKeyRuntimeGeneratedSystemMessage, true) + msgs := []*schema.Message{sys1, sys2} + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 0) + }) + + t.Run("agentic messages - marked stripped", func(t *testing.T) { + sys := schema.SystemAgenticMessage("runtime system") + setMsgExtra(sys, extraKeyRuntimeGeneratedSystemMessage, true) + msgs := []*schema.AgenticMessage{ + sys, + schema.UserAgenticMessage("hello"), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 1) + assert.Equal(t, schema.AgenticRoleTypeUser, result[0].Role) + }) + + t.Run("agentic messages - unmarked preserved", func(t *testing.T) { + sys := schema.SystemAgenticMessage("durable system") + msgs := []*schema.AgenticMessage{ + sys, + schema.UserAgenticMessage("hello"), + } + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 2) + assert.Equal(t, schema.AgenticRoleTypeSystem, result[0].Role) + }) + + t.Run("returns a new slice, not the original slice", func(t *testing.T) { + sys := schema.SystemMessage("runtime") + setMsgExtra(sys, extraKeyRuntimeGeneratedSystemMessage, true) + user := schema.UserMessage("hello") + assistant := schema.AssistantMessage("resp", nil) + msgs := []*schema.Message{sys, user, assistant} + result := stripRuntimeGeneratedLeadingSystemMessages(msgs) + assert.Len(t, result, 2) + originalLen := len(result) + result = append(result, schema.UserMessage("extra")) + assert.Len(t, msgs, 3, "appending to result must not affect original slice") + assert.Len(t, result, originalLen+1) + }) +} + +func TestIsMarkedRuntimeGeneratedSystemMessage(t *testing.T) { + t.Run("marked system message", func(t *testing.T) { + msg := schema.SystemMessage("test") + setMsgExtra(msg, extraKeyRuntimeGeneratedSystemMessage, true) + assert.True(t, isMarkedRuntimeGeneratedSystemMessage(msg)) + }) + + t.Run("unmarked system message", func(t *testing.T) { + msg := schema.SystemMessage("test") + assert.False(t, isMarkedRuntimeGeneratedSystemMessage(msg)) + }) + + t.Run("non-system message with marker", func(t *testing.T) { + msg := schema.UserMessage("test") + setMsgExtra(msg, extraKeyRuntimeGeneratedSystemMessage, true) + assert.False(t, isMarkedRuntimeGeneratedSystemMessage(msg)) + }) + + t.Run("nil extra", func(t *testing.T) { + msg := schema.SystemMessage("test") + msg.Extra = nil + assert.False(t, isMarkedRuntimeGeneratedSystemMessage(msg)) + }) + + t.Run("marker with wrong type", func(t *testing.T) { + msg := schema.SystemMessage("test") + setMsgExtra(msg, extraKeyRuntimeGeneratedSystemMessage, "yes") + assert.False(t, isMarkedRuntimeGeneratedSystemMessage(msg)) + }) +} diff --git a/adk/prebuilt/deep/checkpoint_compat_resume_test.go b/adk/prebuilt/deep/checkpoint_compat_resume_test.go index 1a4f8baa7..744549ee6 100644 --- a/adk/prebuilt/deep/checkpoint_compat_resume_test.go +++ b/adk/prebuilt/deep/checkpoint_compat_resume_test.go @@ -172,31 +172,44 @@ func TestDeepAgentCheckpointCompat_V0_8_Resume(t *testing.T) { name string checkpointID string filename string + // brokenByAgentToolInterruptStateChange marks fixtures that were captured + // before the AgentTool interrupt state format was changed to wrap the + // bridge checkpoint bytes inside a JSON envelope (agentToolInterruptState) + // to carry the synthetic child SessionID. The change is documented as + // backward-incompatible in the session event-log reconstruction plan. + brokenByAgentToolInterruptStateChange bool }{ { - name: "v0.7.37", - checkpointID: "checkpoint_compat_v0_7_37", - filename: "checkpoint_data_v0.7.37.bin", + name: "v0.7.37", + checkpointID: "checkpoint_compat_v0_7_37", + filename: "checkpoint_data_v0.7.37.bin", + brokenByAgentToolInterruptStateChange: true, }, { - name: "v0.8.2", - checkpointID: "checkpoint_compat_v0_8_2", - filename: "checkpoint_data_v0.8.2.bin", + name: "v0.8.2", + checkpointID: "checkpoint_compat_v0_8_2", + filename: "checkpoint_data_v0.8.2.bin", + brokenByAgentToolInterruptStateChange: true, }, { - name: "v0.8.3", - checkpointID: "checkpoint_compat_v0_8_3", - filename: "checkpoint_data_v0.8.3.bin", + name: "v0.8.3", + checkpointID: "checkpoint_compat_v0_8_3", + filename: "checkpoint_data_v0.8.3.bin", + brokenByAgentToolInterruptStateChange: true, }, { - name: "v0.8.4", - checkpointID: "checkpoint_compat_v0_8_4", - filename: "checkpoint_data_v0.8.4.bin", + name: "v0.8.4", + checkpointID: "checkpoint_compat_v0_8_4", + filename: "checkpoint_data_v0.8.4.bin", + brokenByAgentToolInterruptStateChange: true, }, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { + if tc.brokenByAgentToolInterruptStateChange { + t.Skip("AgentTool interrupt state format changed for SessionID-based event filtering; pre-change checkpoint fixtures are not resumable. See plan-session-event-log-reconstruction.md.") + } runDeepAgentCheckpointCompat(t, tc.checkpointID, tc.filename) }) } diff --git a/adk/prebuilt/deep/deep.go b/adk/prebuilt/deep/deep.go index 531511319..b91ee8349 100644 --- a/adk/prebuilt/deep/deep.go +++ b/adk/prebuilt/deep/deep.go @@ -1,5 +1,5 @@ /* - * Copyright 2025 CloudWeGo Authors + * Copyright 2026 CloudWeGo Authors * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,9 +24,12 @@ import ( "github.com/bytedance/sonic" "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/backgroundtask" "github.com/cloudwego/eino/adk/filesystem" "github.com/cloudwego/eino/adk/internal" + backgroundtaskmw "github.com/cloudwego/eino/adk/middlewares/backgroundtask" filesystem2 "github.com/cloudwego/eino/adk/middlewares/filesystem" + "github.com/cloudwego/eino/adk/middlewares/subagent" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/components/tool/utils" "github.com/cloudwego/eino/schema" @@ -37,6 +40,25 @@ func init() { schema.RegisterName[[]TODO]("_eino_adk_prebuilt_deep_todo_slice") } +// BackgroundConfig enables background-task execution for a DeepAgent's top-level +// agent. When set, shell commands and sub-agent runs can execute as managed +// background tasks under one task-ID space, and the task_output/task_stop control +// tools are injected once. +type BackgroundConfig struct { + // Manager is the shared background-task Manager. Required (a nil Manager is the + // same as no BackgroundConfig). + Manager *backgroundtask.Manager + + // OutputDir, when set together with Config.Backend, gives every managed + // background task (shell command or sub-agent run) an output file under this + // directory. Shell runs tee their output there as it streams (interim output); + // sub-agent runs write their final result there. The path is recorded on + // Task.OutputFile and surfaced when the task is launched in the background, so a + // backgrounded task's output is retrievable by path. When empty, tasks have no + // output file. + OutputDir string +} + // TypedConfig defines the configuration for creating a DeepAgent parameterized by message type. // An Agentic DeepAgent (M = *schema.AgenticMessage) only supports Agentic sub-agents, // and a standard DeepAgent (M = *schema.Message) only supports standard sub-agents. @@ -65,17 +87,32 @@ type TypedConfig[M adk.MessageType] struct { // Backend provides filesystem operations used by tools and offloading. // If set, filesystem tools (read_file, write_file, edit_file, glob, grep) will be registered. + // For advanced filesystem middleware configuration, leave Backend, Shell, and StreamingShell empty + // and pass a manually constructed filesystem middleware through Handlers. // Optional. Backend filesystem.Backend // Shell provides shell command execution capability. // If set, an execute tool will be registered to support shell command execution. + // For advanced filesystem middleware configuration, leave Backend, Shell, and StreamingShell empty + // and pass a manually constructed filesystem middleware through Handlers. // Optional. Mutually exclusive with StreamingShell. Shell filesystem.Shell // StreamingShell provides streaming shell command execution capability. // If set, a streaming execute tool will be registered to support streaming shell command execution. + // For advanced filesystem middleware configuration, leave Backend, Shell, and StreamingShell empty + // and pass a manually constructed filesystem middleware through Handlers. // Optional. Mutually exclusive with Shell. StreamingShell filesystem.StreamingShell + // Background configures background-task execution for the top-level agent: it + // can spawn sub-agents and run shell commands as managed background tasks under + // one task-ID space, and the task_output/task_stop control tools are injected + // once. Background is intentionally NOT propagated to the general or user + // sub-agents: their shell runs stay foreground/buffered and they cannot launch + // background work, so background orchestration is a top-level concern only. When + // nil, the top-level agent has no background-task support. See BackgroundConfig. + Background *BackgroundConfig + // WithoutWriteTodos disables the built-in write_todos tool when set to true. WithoutWriteTodos bool // WithoutGeneralSubAgent disables the general-purpose subagent when set to true. @@ -101,7 +138,6 @@ type TypedConfig[M adk.MessageType] struct { // When set, the agent will automatically fail over to alternative models on errors. // This config is also propagated to the general sub-agent. ModelFailoverConfig *adk.ModelFailoverConfig[M] - // OutputKey stores the agent's response in the session. // Optional. When set, stores output via AddSessionValue(ctx, outputKey, msg.Content). OutputKey string @@ -114,7 +150,9 @@ type Config = TypedConfig[*schema.Message] // This function initializes built-in tools, creates a task tool for subagent orchestration, // and returns a fully configured TypedChatModelAgent ready for execution. func NewTyped[M adk.MessageType](ctx context.Context, cfg *TypedConfig[M]) (adk.TypedResumableAgent[M], error) { - handlers, err := buildTypedBuiltinAgentMiddlewares(ctx, cfg) + // Sub-agents never get the Manager: their shell runs stay foreground/buffered + // and they cannot launch background work (see Config.Manager). + subAgentHandlers, err := buildTypedBuiltinAgentMiddlewares(ctx, cfg, nil) if err != nil { return nil, err } @@ -127,25 +165,47 @@ func NewTyped[M adk.MessageType](ctx context.Context, cfg *TypedConfig[M]) (adk. }) } + // The top-level agent's built-in handlers do get background support, so its own + // shell runs are background-capable and tracked under the shared task-ID space. + handlers, err := buildTypedBuiltinAgentMiddlewares(ctx, cfg, cfg.Background) + if err != nil { + return nil, err + } + if !cfg.WithoutGeneralSubAgent || len(cfg.SubAgents) > 0 { - tt, err := typedTaskToolMiddleware( - ctx, - cfg.TaskToolDescriptionGenerator, - cfg.SubAgents, - - cfg.WithoutGeneralSubAgent, - cfg.ChatModel, - instruction, - cfg.ToolsConfig, - cfg.MaxIteration, - cfg.Middlewares, - append(handlers, cfg.Handlers...), - cfg.ModelFailoverConfig, - ) + allSubAgents, err := buildSubAgentsList(ctx, cfg, instruction, subAgentHandlers) + if err != nil { + return nil, err + } + if len(allSubAgents) > 0 { + subCfg := &subagent.TypedConfig[M]{ + SubAgents: allSubAgents, + ToolName: taskToolName, + ToolDescriptionGenerator: cfg.TaskToolDescriptionGenerator, + } + if cfg.Background != nil && cfg.Background.Manager != nil { + subCfg.Background = &subagent.BackgroundConfig{ + Manager: cfg.Background.Manager, + OutputStore: backendAppender(cfg.Backend), + OutputDir: cfg.Background.OutputDir, + } + } + subagentMW, err := subagent.NewTyped[M](ctx, subCfg) + if err != nil { + return nil, fmt.Errorf("failed to create subagent middleware: %w", err) + } + handlers = append(handlers, subagentMW) + } + } + + // When background support is configured, wire its control tools + // (task_output/task_stop) exactly once at the top level. + if cfg.Background != nil && cfg.Background.Manager != nil { + controlMW, err := backgroundtaskmw.NewTyped[M](ctx, &backgroundtaskmw.TypedConfig[M]{Manager: cfg.Background.Manager}) if err != nil { - return nil, fmt.Errorf("failed to new task tool: %w", err) + return nil, fmt.Errorf("failed to create background-task control middleware: %w", err) } - handlers = append(handlers, tt) + handlers = append(handlers, controlMW) } return adk.NewTypedChatModelAgent(ctx, &adk.TypedChatModelAgentConfig[M]{ @@ -177,11 +237,17 @@ func typedGenModelInput[M adk.MessageType](_ context.Context, instruction string switch any(zero).(type) { case *schema.Message: msgs := make([]*schema.Message, 0, len(input.Messages)+1) + inputMessages := input.Messages if instruction != "" { + if len(inputMessages) > 0 { + if msg, ok := any(inputMessages[0]).(*schema.Message); ok && msg.Role == schema.System { + inputMessages = inputMessages[1:] + } + } msgs = append(msgs, schema.SystemMessage(instruction)) } // Type assertion is safe here because M = *schema.Message. - for _, m := range input.Messages { + for _, m := range inputMessages { msgs = append(msgs, any(m).(*schema.Message)) } result := make([]M, len(msgs)) @@ -191,10 +257,16 @@ func typedGenModelInput[M adk.MessageType](_ context.Context, instruction string return result, nil case *schema.AgenticMessage: msgs := make([]*schema.AgenticMessage, 0, len(input.Messages)+1) + inputMessages := input.Messages if instruction != "" { + if len(inputMessages) > 0 { + if msg, ok := any(inputMessages[0]).(*schema.AgenticMessage); ok && msg.Role == schema.AgenticRoleTypeSystem { + inputMessages = inputMessages[1:] + } + } msgs = append(msgs, schema.SystemAgenticMessage(instruction)) } - for _, m := range input.Messages { + for _, m := range inputMessages { msgs = append(msgs, any(m).(*schema.AgenticMessage)) } result := make([]M, len(msgs)) @@ -206,7 +278,38 @@ func typedGenModelInput[M adk.MessageType](_ context.Context, instruction string panic("unreachable") } -func buildTypedBuiltinAgentMiddlewares[M adk.MessageType](ctx context.Context, cfg *TypedConfig[M]) ([]adk.TypedChatModelAgentMiddleware[M], error) { +func buildSubAgentsList[M adk.MessageType](ctx context.Context, cfg *TypedConfig[M], instruction string, handlers []adk.TypedChatModelAgentMiddleware[M]) ([]adk.TypedAgent[M], error) { + var allSubAgents []adk.TypedAgent[M] + + if !cfg.WithoutGeneralSubAgent { + agentDesc := internal.SelectPrompt(internal.I18nPrompts{ + English: generalAgentDescription, + Chinese: generalAgentDescriptionChinese, + }) + generalAgent, err := adk.NewTypedChatModelAgent(ctx, &adk.TypedChatModelAgentConfig[M]{ + Name: generalAgentName, + Description: agentDesc, + Instruction: instruction, + Model: cfg.ChatModel, + ToolsConfig: cfg.ToolsConfig, + MaxIterations: cfg.MaxIteration, + Middlewares: cfg.Middlewares, + Handlers: append(handlers, cfg.Handlers...), + GenModelInput: typedGenModelInput[M], + ModelRetryConfig: cfg.ModelRetryConfig, + ModelFailoverConfig: cfg.ModelFailoverConfig, + }) + if err != nil { + return nil, err + } + allSubAgents = append(allSubAgents, generalAgent) + } + + allSubAgents = append(allSubAgents, cfg.SubAgents...) + return allSubAgents, nil +} + +func buildTypedBuiltinAgentMiddlewares[M adk.MessageType](ctx context.Context, cfg *TypedConfig[M], background *BackgroundConfig) ([]adk.TypedChatModelAgentMiddleware[M], error) { var ms []adk.TypedChatModelAgentMiddleware[M] if !cfg.WithoutWriteTodos { t, err := typedNewWriteTodos[M]() @@ -217,11 +320,19 @@ func buildTypedBuiltinAgentMiddlewares[M adk.MessageType](ctx context.Context, c } if cfg.Backend != nil || cfg.Shell != nil || cfg.StreamingShell != nil { - fm, err := filesystem2.NewTyped[M](ctx, &filesystem2.MiddlewareConfig{ + mwCfg := &filesystem2.MiddlewareConfig{ Backend: cfg.Backend, Shell: cfg.Shell, StreamingShell: cfg.StreamingShell, - }) + } + if background != nil && background.Manager != nil { + mwCfg.Background = &filesystem2.BackgroundConfig{ + Manager: background.Manager, + OutputStore: backendAppender(cfg.Backend), + OutputDir: background.OutputDir, + } + } + fm, err := filesystem2.NewTyped[M](ctx, mwCfg) if err != nil { return nil, err } @@ -231,6 +342,14 @@ func buildTypedBuiltinAgentMiddlewares[M adk.MessageType](ctx context.Context, c return ms, nil } +// backendAppender returns b as a filesystem.Appender when it supports incremental +// append, or nil otherwise — in which case background tasks run without output +// files. The default InMemoryBackend implements Appender. +func backendAppender(b filesystem.Backend) filesystem.Appender { + ap, _ := b.(filesystem.Appender) + return ap +} + type TODO struct { Content string `json:"content"` ActiveForm string `json:"activeForm"` diff --git a/adk/prebuilt/deep/deep_test.go b/adk/prebuilt/deep/deep_test.go index b39cfe9f5..93fedc311 100644 --- a/adk/prebuilt/deep/deep_test.go +++ b/adk/prebuilt/deep/deep_test.go @@ -1,5 +1,5 @@ /* - * Copyright 2025 CloudWeGo Authors + * Copyright 2026 CloudWeGo Authors * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -20,6 +20,7 @@ import ( "context" "fmt" "io" + "sync" "sync/atomic" "testing" @@ -27,7 +28,11 @@ import ( "go.uber.org/mock/gomock" "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/backgroundtask" + "github.com/cloudwego/eino/adk/filesystem" + filesystem2 "github.com/cloudwego/eino/adk/middlewares/filesystem" "github.com/cloudwego/eino/adk/prebuilt/planexecute" + adksession "github.com/cloudwego/eino/adk/session" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/compose" @@ -89,6 +94,56 @@ func readAgenticText(msg *schema.AgenticMessage) string { return "" } +func readAgenticInputText(msg *schema.AgenticMessage) string { + if msg == nil { + return "" + } + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + if block.UserInputText != nil { + return block.UserInputText.Text + } + } + return "" +} + +type recordingDeepModel struct { + mu sync.Mutex + inputs [][]*schema.Message + response *schema.Message +} + +func (m *recordingDeepModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.mu.Lock() + defer m.mu.Unlock() + copied := append([]*schema.Message{}, input...) + m.inputs = append(m.inputs, copied) + if m.response != nil { + return m.response, nil + } + return schema.AssistantMessage("ok", nil), nil +} + +func (m *recordingDeepModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *recordingDeepModel) snapshotInputs() [][]*schema.Message { + m.mu.Lock() + defer m.mu.Unlock() + out := make([][]*schema.Message, len(m.inputs)) + for i, input := range m.inputs { + out[i] = append([]*schema.Message{}, input...) + } + return out +} + type mockSearchTool struct{} func (m *mockSearchTool) Info(context.Context) (*schema.ToolInfo, error) { @@ -102,6 +157,12 @@ func (m *mockSearchTool) InvokableRun(context.Context, string, ...tool.Option) ( return "latest news search result", nil } +type deepMockShell struct{} + +func (m *deepMockShell) Execute(ctx context.Context, req *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) { + return &filesystem.ExecuteResponse{Output: "ok"}, nil +} + func TestGenModelInput(t *testing.T) { ctx := context.Background() @@ -121,6 +182,56 @@ func TestGenModelInput(t *testing.T) { assert.Equal(t, "hello", msgs[1].Content) }) + t.Run("WithInstructionStripsLeadingSystemMessage", func(t *testing.T) { + input := &adk.AgentInput{ + Messages: []*schema.Message{ + schema.SystemMessage("old"), + schema.UserMessage("hello"), + }, + } + + msgs, err := typedGenModelInput(ctx, "new", input) + assert.NoError(t, err) + assert.Len(t, msgs, 2) + assert.Equal(t, schema.System, msgs[0].Role) + assert.Equal(t, "new", msgs[0].Content) + assert.Equal(t, schema.User, msgs[1].Role) + assert.Equal(t, "hello", msgs[1].Content) + + systemCount := 0 + for _, msg := range msgs { + if msg.Role == schema.System { + systemCount++ + } + } + assert.Equal(t, 1, systemCount) + }) + + t.Run("WithInstructionStripsLeadingAgenticSystemMessage", func(t *testing.T) { + input := &adk.TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.SystemAgenticMessage("old"), + schema.UserAgenticMessage("hello"), + }, + } + + msgs, err := typedGenModelInput(ctx, "new", input) + assert.NoError(t, err) + assert.Len(t, msgs, 2) + assert.Equal(t, schema.AgenticRoleTypeSystem, msgs[0].Role) + assert.Equal(t, "new", readAgenticInputText(msgs[0])) + assert.Equal(t, schema.AgenticRoleTypeUser, msgs[1].Role) + assert.Equal(t, "hello", readAgenticInputText(msgs[1])) + + systemCount := 0 + for _, msg := range msgs { + if msg.Role == schema.AgenticRoleTypeSystem { + systemCount++ + } + } + assert.Equal(t, 1, systemCount) + }) + t.Run("WithoutInstruction", func(t *testing.T) { input := &adk.AgentInput{ Messages: []*schema.Message{ @@ -134,10 +245,87 @@ func TestGenModelInput(t *testing.T) { assert.Equal(t, schema.User, msgs[0].Role) assert.Equal(t, "hello", msgs[0].Content) }) + + t.Run("WithoutInstructionPreservesLeadingSystemMessage", func(t *testing.T) { + input := &adk.AgentInput{ + Messages: []*schema.Message{ + schema.SystemMessage("old"), + schema.UserMessage("hello"), + }, + } + + msgs, err := typedGenModelInput(ctx, "", input) + assert.NoError(t, err) + assert.Len(t, msgs, 2) + assert.Equal(t, schema.System, msgs[0].Role) + assert.Equal(t, "old", msgs[0].Content) + assert.Equal(t, schema.User, msgs[1].Role) + assert.Equal(t, "hello", msgs[1].Content) + }) +} + +func TestDeepAgentTurn2DeduplicatesPersistedLeadingSystemMessage(t *testing.T) { + ctx := context.Background() + store := adksession.NewInMemoryStore[*schema.Message](nil) + model := &recordingDeepModel{} + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: model, + Instruction: "you are deep agent", + MaxIteration: 2, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + assert.NoError(t, err) + if err != nil { + return + } + + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + SessionID: "deep-leading-system-dedup", + SessionStore: store, + }) + for _, input := range [][]adk.Message{ + {schema.UserMessage("turn one")}, + {schema.UserMessage("turn two")}, + } { + iter := runner.Run(ctx, input) + for { + if _, ok := iter.Next(); !ok { + break + } + } + } + + inputs := model.snapshotInputs() + assert.Len(t, inputs, 2) + if len(inputs) < 2 { + return + } + secondTurnInput := inputs[1] + assert.NotEmpty(t, secondTurnInput) + if len(secondTurnInput) == 0 { + return + } + assert.Equal(t, schema.System, secondTurnInput[0].Role) + assert.Equal(t, "you are deep agent", secondTurnInput[0].Content) + + systemCount := 0 + for _, msg := range secondTurnInput { + if msg.Role == schema.System { + systemCount++ + } + } + assert.Equal(t, 1, systemCount) + if len(secondTurnInput) > 1 { + assert.NotEqual(t, schema.System, secondTurnInput[1].Role, "reconstructed system message must not remain after fresh system message") + } } func TestWriteTodos(t *testing.T) { - m, err := buildTypedBuiltinAgentMiddlewares(context.Background(), &Config{WithoutWriteTodos: false}) + m, err := buildTypedBuiltinAgentMiddlewares(context.Background(), &Config{WithoutWriteTodos: false}, nil) assert.NoError(t, err) wt := m[0].(*typedAppendPromptTool[*schema.Message]).t.(tool.InvokableTool) @@ -150,6 +338,176 @@ func TestWriteTodos(t *testing.T) { assert.Equal(t, fmt.Sprintf("Updated todo list to %s", todos), result) } +func TestDeepAgentFilesystemExecuteDefaults(t *testing.T) { + ctx := context.Background() + backend := filesystem.NewInMemoryBackend() + + tests := []struct { + name string + cfg *Config + wantToolLen int + }{ + { + name: "backend and shell", + cfg: &Config{ + WithoutWriteTodos: true, + Backend: backend, + Shell: &deepMockShell{}, + }, + wantToolLen: 7, + }, + { + name: "shell only", + cfg: &Config{ + WithoutWriteTodos: true, + Shell: &deepMockShell{}, + }, + wantToolLen: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + handlers, err := buildTypedBuiltinAgentMiddlewares(ctx, tt.cfg, nil) + assert.NoError(t, err) + assert.Len(t, handlers, 1) + + _, runCtx, err := handlers[0].BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{}) + assert.NoError(t, err) + assert.NotNil(t, runCtx) + assert.Len(t, runCtx.Tools, tt.wantToolLen) + + toolNames := make(map[string]bool) + var executeTool tool.BaseTool + for _, tl := range runCtx.Tools { + info, infoErr := tl.Info(ctx) + assert.NoError(t, infoErr) + toolNames[info.Name] = true + if info.Name == filesystem2.ToolNameExecute { + executeTool = tl + } + } + + assert.NotNil(t, executeTool) + assert.False(t, toolNames["execute_output"]) + assert.False(t, toolNames["execute_wait"]) + assert.False(t, toolNames["execute_stop"]) + assert.False(t, toolNames["execute_list"]) + + info, err := executeTool.Info(ctx) + assert.NoError(t, err) + js, err := info.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + _, ok := js.Properties.Get("command") + assert.True(t, ok) + _, ok = js.Properties.Get("mode") + assert.False(t, ok) + _, ok = js.Properties.Get("wait_ms") + assert.False(t, ok) + }) + } +} + +func TestDeepAgentManagerWiring(t *testing.T) { + ctx := context.Background() + + // With a Manager, the top-level built-in handlers route execute through it, so + // the execute tool gains a run_in_background field. + mgr := backgroundtask.New(ctx, &backgroundtask.Config{}) + defer func() { _ = mgr.Close(ctx) }() + + handlers, err := buildTypedBuiltinAgentMiddlewares(ctx, &Config{ + WithoutWriteTodos: true, + Shell: &deepMockShell{}, + }, &BackgroundConfig{Manager: mgr}) + assert.NoError(t, err) + assert.Len(t, handlers, 1) + + _, runCtx, err := handlers[0].BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{}) + assert.NoError(t, err) + assert.NotNil(t, runCtx) + assert.Len(t, runCtx.Tools, 1) + + info, err := runCtx.Tools[0].Info(ctx) + assert.NoError(t, err) + js, err := info.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + _, ok := js.Properties.Get("run_in_background") + assert.True(t, ok, "managed execute must expose run_in_background") + + // Without a Manager, the same handlers produce a command-only execute tool. + plain, err := buildTypedBuiltinAgentMiddlewares(ctx, &Config{ + WithoutWriteTodos: true, + Shell: &deepMockShell{}, + }, nil) + assert.NoError(t, err) + _, plainCtx, err := plain[0].BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{}) + assert.NoError(t, err) + plainInfo, err := plainCtx.Tools[0].Info(ctx) + assert.NoError(t, err) + plainJS, err := plainInfo.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + _, ok = plainJS.Properties.Get("run_in_background") + assert.False(t, ok, "unmanaged execute must not expose run_in_background") +} + +// NewTyped with a Manager injects the task_output/task_stop control tools and a +// background-capable subagent tool exactly once at the top level. +func TestDeepAgentNewTypedWithManager(t *testing.T) { + ctx := context.Background() + mgr := backgroundtask.New(ctx, &backgroundtask.Config{}) + defer func() { _ = mgr.Close(ctx) }() + + cm := mockModel.NewMockToolCallingChatModel(gomock.NewController(t)) + + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: cm, + Shell: &deepMockShell{}, + Background: &BackgroundConfig{Manager: mgr}, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) +} + +func TestDeepAgentManualFilesystemMiddlewarePath(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + fsMW, err := filesystem2.New(ctx, &filesystem2.MiddlewareConfig{ + Shell: &deepMockShell{}, + ExecuteToolConfig: &filesystem2.ExecuteToolConfig{}, + }) + assert.NoError(t, err) + + _, runCtx, err := fsMW.BeforeAgent(ctx, &adk.ChatModelAgentContext[*schema.Message]{}) + assert.NoError(t, err) + assert.Len(t, runCtx.Tools, 1) + info, err := runCtx.Tools[0].Info(ctx) + assert.NoError(t, err) + assert.Equal(t, filesystem2.ToolNameExecute, info.Name) + js, err := info.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + _, ok := js.Properties.Get("command") + assert.True(t, ok) + + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: cm, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + Handlers: []adk.ChatModelAgentMiddleware{fsMW}, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) +} + func TestDeepSubAgentSharesSessionValues(t *testing.T) { ctx := context.Background() spy := &spySubAgent{} diff --git a/adk/prebuilt/deep/task_tool.go b/adk/prebuilt/deep/task_tool.go deleted file mode 100644 index 5c7e50b63..000000000 --- a/adk/prebuilt/deep/task_tool.go +++ /dev/null @@ -1,191 +0,0 @@ -/* - * Copyright 2025 CloudWeGo Authors - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package deep - -import ( - "context" - "encoding/json" - "fmt" - "strings" - - "github.com/bytedance/sonic" - "github.com/slongfield/pyfmt" - - "github.com/cloudwego/eino/adk" - "github.com/cloudwego/eino/adk/internal" - "github.com/cloudwego/eino/components/model" - "github.com/cloudwego/eino/components/tool" - "github.com/cloudwego/eino/schema" -) - -func typedTaskToolMiddleware[M adk.MessageType]( - ctx context.Context, - taskToolDescriptionGenerator func(ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error), - subAgents []adk.TypedAgent[M], - - withoutGeneralSubAgent bool, - cm model.BaseModel[M], - instruction string, - toolsConfig adk.ToolsConfig, - maxIteration int, - middlewares []adk.AgentMiddleware, - handlers []adk.TypedChatModelAgentMiddleware[M], - modelFailoverConfig *adk.ModelFailoverConfig[M], -) (adk.TypedChatModelAgentMiddleware[M], error) { - t, err := typedNewTaskTool(ctx, taskToolDescriptionGenerator, subAgents, withoutGeneralSubAgent, cm, instruction, toolsConfig, maxIteration, middlewares, handlers, modelFailoverConfig) - if err != nil { - return nil, err - } - prompt := internal.SelectPrompt(internal.I18nPrompts{ - English: taskPrompt, - Chinese: taskPromptChinese, - }) - - return typedBuildAppendPromptTool[M](prompt, t), nil -} - -func typedNewTaskTool[M adk.MessageType]( - ctx context.Context, - taskToolDescriptionGenerator func(ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error), - subAgents []adk.TypedAgent[M], - - withoutGeneralSubAgent bool, - cm model.BaseModel[M], - instruction string, - toolsConfig adk.ToolsConfig, - maxIteration int, - middlewares []adk.AgentMiddleware, - handlers []adk.TypedChatModelAgentMiddleware[M], - modelFailoverConfig *adk.ModelFailoverConfig[M], -) (tool.InvokableTool, error) { - t := &typedTaskTool[M]{ - subAgents: map[string]tool.InvokableTool{}, - subAgentSlice: subAgents, - descGen: typedDefaultTaskToolDescription[M], - } - - if taskToolDescriptionGenerator != nil { - t.descGen = taskToolDescriptionGenerator - } - - if !withoutGeneralSubAgent { - agentDesc := internal.SelectPrompt(internal.I18nPrompts{ - English: generalAgentDescription, - Chinese: generalAgentDescriptionChinese, - }) - generalAgent, err := adk.NewTypedChatModelAgent(ctx, &adk.TypedChatModelAgentConfig[M]{ - Name: generalAgentName, - Description: agentDesc, - Instruction: instruction, - Model: cm, - ToolsConfig: toolsConfig, - MaxIterations: maxIteration, - Middlewares: middlewares, - Handlers: handlers, - GenModelInput: typedGenModelInput[M], - ModelFailoverConfig: modelFailoverConfig, - }) - if err != nil { - return nil, err - } - - it, err := assertAgentTool(adk.NewTypedAgentTool(ctx, adk.TypedAgent[M](generalAgent))) - if err != nil { - return nil, err - } - t.subAgents[generalAgent.Name(ctx)] = it - t.subAgentSlice = append(t.subAgentSlice, generalAgent) - } - - for _, a := range subAgents { - name := a.Name(ctx) - it, err := assertAgentTool(adk.NewTypedAgentTool(ctx, a)) - if err != nil { - return nil, err - } - t.subAgents[name] = it - } - - return t, nil -} - -type typedTaskTool[M adk.MessageType] struct { - subAgents map[string]tool.InvokableTool - subAgentSlice []adk.TypedAgent[M] - descGen func(ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error) -} - -func (t *typedTaskTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) { - desc, err := t.descGen(ctx, t.subAgentSlice) - if err != nil { - return nil, err - } - return &schema.ToolInfo{ - Name: taskToolName, - Desc: desc, - ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ - "subagent_type": { - Type: schema.String, - }, - "description": { - Type: schema.String, - }, - }), - }, nil -} - -type taskToolArgument struct { - SubagentType string `json:"subagent_type"` - Description string `json:"description"` -} - -func (t *typedTaskTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { - input := &taskToolArgument{} - err := json.Unmarshal([]byte(argumentsInJSON), input) - if err != nil { - return "", fmt.Errorf("failed to unmarshal task tool input json: %w", err) - } - a, ok := t.subAgents[input.SubagentType] - if !ok { - return "", fmt.Errorf("subagent type %s not found", input.SubagentType) - } - - params, err := sonic.MarshalString(map[string]string{ - "request": input.Description, - }) - if err != nil { - return "", err - } - - return a.InvokableRun(ctx, params, opts...) -} - -func typedDefaultTaskToolDescription[M adk.MessageType](ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error) { - subAgentsDescBuilder := strings.Builder{} - for _, a := range subAgents { - name := a.Name(ctx) - desc := a.Description(ctx) - subAgentsDescBuilder.WriteString(fmt.Sprintf("- %s: %s\n", name, desc)) - } - toolDesc := internal.SelectPrompt(internal.I18nPrompts{ - English: taskToolDescription, - Chinese: taskToolDescriptionChinese, - }) - return pyfmt.Fmt(toolDesc, map[string]any{ - "other_agents": subAgentsDescBuilder.String(), - }) -} diff --git a/adk/prebuilt/deep/task_tool_test.go b/adk/prebuilt/deep/task_tool_test.go deleted file mode 100644 index 44fa80b8e..000000000 --- a/adk/prebuilt/deep/task_tool_test.go +++ /dev/null @@ -1,78 +0,0 @@ -/* - * Copyright 2025 CloudWeGo Authors - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package deep - -import ( - "context" - "testing" - - "github.com/stretchr/testify/assert" - - "github.com/cloudwego/eino/adk" - "github.com/cloudwego/eino/schema" -) - -func TestTaskTool(t *testing.T) { - a1 := &myAgent{name: "1", desc: "desc of my agent 1"} - a2 := &myAgent{name: "2", desc: "desc of my agent 2"} - ctx := context.Background() - tt, err := typedNewTaskTool( - ctx, - nil, - []adk.Agent{a1, a2}, - true, - nil, - "", - adk.ToolsConfig{}, - 10, - nil, - nil, - nil, - ) - assert.NoError(t, err) - - info, err := tt.Info(ctx) - assert.NoError(t, err) - assert.Contains(t, info.Desc, "desc of my agent 1") - - result, err := tt.InvokableRun(ctx, `{"subagent_type":"1"}`) - assert.NoError(t, err) - assert.Equal(t, "desc of my agent 1", result) - result, err = tt.InvokableRun(ctx, `{"subagent_type":"2"}`) - assert.NoError(t, err) - assert.Equal(t, "desc of my agent 2", result) -} - -type myAgent struct { - name string - desc string -} - -func (m *myAgent) Name(_ context.Context) string { - return m.name -} - -func (m *myAgent) Description(_ context.Context) string { - return m.desc -} - -func (m *myAgent) Run(_ context.Context, _ *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { - iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() - gen.Send(adk.EventFromMessage(schema.UserMessage(m.desc), nil, schema.User, "")) - gen.Close() - return iter -} diff --git a/adk/prebuilt/deep/types.go b/adk/prebuilt/deep/types.go index 781418bf3..b2798b2a8 100644 --- a/adk/prebuilt/deep/types.go +++ b/adk/prebuilt/deep/types.go @@ -18,7 +18,6 @@ package deep import ( "context" - "fmt" "github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/components/tool" @@ -33,14 +32,6 @@ const ( SessionKeyTodos = "deep_agent_session_key_todos" ) -func assertAgentTool(t tool.BaseTool) (tool.InvokableTool, error) { - it, ok := t.(tool.InvokableTool) - if !ok { - return nil, fmt.Errorf("failed to assert agent tool type: %T", t) - } - return it, nil -} - func typedBuildAppendPromptTool[M adk.MessageType](prompt string, t tool.BaseTool) adk.TypedChatModelAgentMiddleware[M] { return &typedAppendPromptTool[M]{ TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[M]{}, @@ -55,7 +46,7 @@ type typedAppendPromptTool[M adk.MessageType] struct { prompt string } -func (w *typedAppendPromptTool[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { +func (w *typedAppendPromptTool[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext[M]) (context.Context, *adk.ChatModelAgentContext[M], error) { nRunCtx := *runCtx nRunCtx.Instruction += w.prompt if w.t != nil { diff --git a/adk/react.go b/adk/react.go index 03ef04565..5bb77b717 100644 --- a/adk/react.go +++ b/adk/react.go @@ -22,6 +22,7 @@ import ( "encoding/gob" "errors" "io" + "time" "github.com/cloudwego/eino/adk/internal" "github.com/cloudwego/eino/components/model" @@ -54,6 +55,49 @@ type typedState[M MessageType] struct { ReturnDirectlyEvent *TypedAgentEvent[M] RetryAttempt int ToolMsgIDs map[string]map[string]string // toolName → callID → eino message ID + + // CurrentModelSpanID is the SpanID of the model request span that emitted + // the most recent assistant message containing tool calls. The tool wrapper + // snapshots this value into ToolSpansInFlight when emitting a tool_call_start + // span; the snapshot survives interrupt/resume so the matching tool_call_end + // span (which may be emitted on a later run) preserves the link. + CurrentModelSpanID string + + // CurrentAssistantMessageEventID is the SessionEvent EventID of the most + // recent assistant message that emitted tool calls. Snapshotted into + // ToolSpansInFlight at start emission time, same lifecycle as CurrentModelSpanID. + CurrentAssistantMessageEventID string + + // ToolSpansInFlight tracks tool calls whose tool_call_start span has been + // emitted but whose tool_call_end span has not yet fired (typically because + // the call is paused on an interrupt awaiting user resume). Keyed by + // tCtx.CallID. Entries are inserted at start emission, retained across + // interrupt boundaries, and deleted when the matching end span fires. + ToolSpansInFlight map[string]*toolSpanInFlight +} + +// toolSpanInFlight holds identity for a tool_call_start span that has been +// emitted but whose matching tool_call_end span has not yet fired. The +// typical reason is that the call is paused on a permission interrupt +// awaiting user resume. +// +// The wrapper at typedEventSenderToolWrapper persists one entry per +// tCtx.CallID at start emission. On every subsequent invocation of the +// wrapper for the same CallID (i.e. on resume), the entry is reused so +// that the matching end span carries the same SpanID / StartEventID / +// parent IDs — preserving the temporal semantics that one logical tool +// call corresponds to one logical span pair, even when start and end +// straddle an interrupt boundary. +// +// The entry is deleted when the matching end span is emitted (success, +// hard error, or cancellation). It is NOT deleted on interrupt-shape +// errors; those leave the entry intact for the next resume. +type toolSpanInFlight struct { + SpanID string + StartEventID string + StartedAt time.Time + ParentSpanID string + AssistantMessageEventID string } // State is the internal state of the ChatModelAgent. @@ -427,7 +471,7 @@ func newReact(ctx context.Context, config *reactConfig) (reactGraph, error) { } toolPostHandle := func(ctx context.Context, out *schema.StreamReader[[]*schema.Message], st *State) (*schema.StreamReader[[]*schema.Message], error) { if event := st.getReturnDirectlyEvent(); event != nil { - getTypedChatModelAgentExecCtx[*schema.Message](ctx).send(event) + getTypedChatModelAgentExecCtx[*schema.Message](ctx).send(ctx, event) st.setReturnDirectlyEvent(nil) } return out, nil @@ -675,7 +719,7 @@ func newAgenticReact(ctx context.Context, config *agenticReactConfig) (agenticRe } toolPostHandle := func(ctx context.Context, out *schema.StreamReader[[]*schema.AgenticMessage], st *agenticState) (*schema.StreamReader[[]*schema.AgenticMessage], error) { if event := st.getReturnDirectlyEvent(); event != nil { - getTypedChatModelAgentExecCtx[*schema.AgenticMessage](ctx).send(event) + getTypedChatModelAgentExecCtx[*schema.AgenticMessage](ctx).send(ctx, event) st.setReturnDirectlyEvent(nil) } return out, nil diff --git a/adk/retry_chatmodel.go b/adk/retry_chatmodel.go index 350a3c4a6..c97334b9e 100644 --- a/adk/retry_chatmodel.go +++ b/adk/retry_chatmodel.go @@ -255,7 +255,15 @@ type TypedModelRetryConfig[M MessageType] struct { // ModelRetryConfig is the default retry config type using *schema.Message. type ModelRetryConfig = TypedModelRetryConfig[*schema.Message] +type retryableBeforeOutputError interface { + IsModelTimeoutBeforeOutput() bool +} + func defaultIsRetryAble(_ context.Context, err error) bool { + var timeoutErr retryableBeforeOutputError + if errors.As(err, &timeoutErr) { + return timeoutErr.IsModelTimeoutBeforeOutput() + } return err != nil } @@ -292,6 +300,46 @@ func genErrWrapper(ctx context.Context, maxRetries, attempt int, isRetryAbleFunc } } +func timelineErrorMessage(err error, rejectReason any) string { + if rejectReason != nil { + if msg := fmt.Sprint(rejectReason); msg != "" { + return msg + } + } + if err != nil { + return err.Error() + } + return "" +} + +func emitRetryingTimeline[M MessageType](ctx context.Context, err error, rejectReason ...any) { + var reason any + if len(rejectReason) > 0 { + reason = rejectReason[0] + } + sendSessionTimelineEvent(ctx, &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventSessionError, + Error: &SessionErrorEvent{ + Type: SessionErrorTypeModelRetry, + Message: timelineErrorMessage(err, reason), + RetryStatus: &RetryStatus{Type: "retrying"}, + }, + }) +} + +func emitRetryExhaustedTimeline[M MessageType](ctx context.Context, err error) { + sendSessionTimelineEvent(ctx, &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventSessionError, + Error: &SessionErrorEvent{ + Type: SessionErrorTypeModelRetry, + Message: timelineErrorMessage(err, nil), + RetryStatus: &RetryStatus{Type: "exhausted"}, + }, + }) +} + func consumeStreamForError[M any](stream *schema.StreamReader[M]) error { defer stream.Close() for { @@ -367,12 +415,14 @@ func (r *typedRetryModelWrapper[M]) generateLegacy(ctx context.Context, input [] lastErr = err if attempt < r.config.MaxRetries { + emitRetryingTimeline[M](ctx, err) if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil { return zero, err } } } + emitRetryExhaustedTimeline[M](ctx, lastErr) return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} } @@ -444,7 +494,7 @@ func generateWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx co } if execCtx != nil && execCtx.generator != nil && out != nil { event := typedModelOutputEvent(out, nil) - execCtx.send(event) + execCtx.send(ctx, event) } return out, nil } @@ -458,6 +508,7 @@ func generateWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx co break } + emitRetryingTimeline[M](ctx, lastErr, decision.RejectReason) applyDecisionForRetry(¤tInput, ¤tOpts, ctx, decision) delay := decision.Backoff @@ -470,6 +521,7 @@ func generateWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx co } } + emitRetryExhaustedTimeline[M](ctx, lastErr) return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} } @@ -572,6 +624,7 @@ func streamWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx cont lastErr = err if attempt < r.config.MaxRetries { + emitRetryingTimeline[M](ctx, err) applyDecisionForRetry(¤tInput, ¤tOpts, ctx, decision) delay := decision.Backoff if delay == 0 { @@ -642,6 +695,7 @@ func streamWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx cont lastErr = verdictErr if attempt < r.config.MaxRetries { + emitRetryingTimeline[M](ctx, verdictErr, decision.RejectReason) applyDecisionForRetry(¤tInput, ¤tOpts, ctx, decision) delay := decision.Backoff if delay == 0 { @@ -653,6 +707,7 @@ func streamWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx cont } } + emitRetryExhaustedTimeline[M](ctx, lastErr) return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} } @@ -723,6 +778,7 @@ func (r *typedRetryModelWrapper[M]) streamLegacy(ctx context.Context, input []M, } lastErr = err if attempt < r.config.MaxRetries { + emitRetryingTimeline[M](ctx, err) if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil { return nil, err } @@ -749,11 +805,13 @@ func (r *typedRetryModelWrapper[M]) streamLegacy(ctx context.Context, input []M, lastErr = streamErr if attempt < r.config.MaxRetries { + emitRetryingTimeline[M](ctx, streamErr) if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil { return nil, err } } } + emitRetryExhaustedTimeline[M](ctx, lastErr) return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} } diff --git a/adk/runctx.go b/adk/runctx.go index dd42226af..393d80ccb 100644 --- a/adk/runctx.go +++ b/adk/runctx.go @@ -241,7 +241,8 @@ func GetSessionValue(ctx context.Context, key string) (any, bool) { } func (rs *runSession) addEvent(event *AgentEvent) { - wrapper := &agentEventWrapper{AgentEvent: event, TS: time.Now().UnixNano()} + now := time.Now() + wrapper := &agentEventWrapper{AgentEvent: event, TS: now.UnixNano()} // If LaneEvents is not nil, we are in a parallel lane. // Append to the lane's local event slice (lock-free). if rs.LaneEvents != nil { @@ -298,9 +299,10 @@ func addTypedEvent[M MessageType](session *runSession, event *TypedAgentEvent[M] session.addEvent(any(event).(*AgentEvent)) return } + now := time.Now() session.mtx.Lock() defer session.mtx.Unlock() - wrapper := &typedAgentEventWrapper[M]{event: event, TS: time.Now().UnixNano()} + wrapper := &typedAgentEventWrapper[M]{event: event, TS: now.UnixNano()} store, _ := session.TypedEvents.(*[]*typedAgentEventWrapper[M]) if store == nil { s := make([]*typedAgentEventWrapper[M], 0) @@ -369,6 +371,131 @@ func (rc *runContext) deepCopy() *runContext { return copied } +func sanitizeRunContextForSessionCheckpoint[M MessageType](rc *runContext) *runContext { + if rc == nil { + return nil + } + copied := &runContext{ + RootInput: rc.RootInput, + AgenticRootInput: rc.AgenticRootInput, + RunPath: append([]RunStep(nil), rc.RunPath...), + Session: sanitizeRunSessionForSessionCheckpoint[M](rc.Session), + } + return copied +} + +func sanitizeRunSessionForSessionCheckpoint[M MessageType](rs *runSession) *runSession { + if rs == nil { + return nil + } + + copied := &runSession{ + Values: make(map[string]any), + valuesMtx: &sync.Mutex{}, + } + + if rs.valuesMtx != nil { + rs.valuesMtx.Lock() + for k, v := range rs.Values { + copied.Values[k] = v + } + rs.valuesMtx.Unlock() + } else { + for k, v := range rs.Values { + copied.Values[k] = v + } + } + + var events []*agentEventWrapper + var typedEvents any + rs.mtx.Lock() + events = append(events, rs.Events...) + typedEvents = rs.TypedEvents + rs.mtx.Unlock() + + for _, event := range events { + if sanitized := sanitizeAgentEventWrapperForSessionCheckpoint(event); sanitized != nil { + copied.Events = append(copied.Events, sanitized) + } + } + copied.LaneEvents = sanitizeLaneEventsForSessionCheckpoint(rs.LaneEvents) + + if store, ok := typedEvents.(*[]*typedAgentEventWrapper[M]); ok { + if store == nil { + copied.TypedEvents = store + return copied + } + sanitized := make([]*typedAgentEventWrapper[M], 0, len(*store)) + for _, event := range *store { + if copiedEvent := sanitizeTypedAgentEventWrapperForSessionCheckpoint(event); copiedEvent != nil { + sanitized = append(sanitized, copiedEvent) + } + } + copied.TypedEvents = &sanitized + } else { + copied.TypedEvents = typedEvents + } + + return copied +} + +func sanitizeLaneEventsForSessionCheckpoint(le *laneEvents) *laneEvents { + if le == nil { + return nil + } + copied := &laneEvents{ + Parent: sanitizeLaneEventsForSessionCheckpoint(le.Parent), + } + for _, event := range le.Events { + if sanitized := sanitizeAgentEventWrapperForSessionCheckpoint(event); sanitized != nil { + copied.Events = append(copied.Events, sanitized) + } + } + return copied +} + +func sanitizeAgentEventWrapperForSessionCheckpoint(w *agentEventWrapper) *agentEventWrapper { + if w == nil || w.AgentEvent == nil { + return nil + } + + event := *w.AgentEvent + event.RunPath = append([]RunStep(nil), w.AgentEvent.RunPath...) + event.SessionEventVariant = nil + if event.Output == nil && event.Action == nil && event.Err == nil { + return nil + } + + return &agentEventWrapper{ + AgentEvent: &event, + concatenatedMessage: w.concatenatedMessage, + TS: w.TS, + StreamErr: w.StreamErr, + } +} + +func sanitizeTypedAgentEventWrapperForSessionCheckpoint[M MessageType]( + w *typedAgentEventWrapper[M], +) *typedAgentEventWrapper[M] { + if w == nil || w.event == nil { + return nil + } + + event := *w.event + event.RunPath = append([]RunStep(nil), w.event.RunPath...) + event.SessionEventVariant = nil + if event.Output == nil && event.Action == nil && event.Err == nil { + return nil + } + + return &typedAgentEventWrapper[M]{ + event: &event, + concatenatedMessage: w.concatenatedMessage, + TS: w.TS, + StreamErr: w.StreamErr, + } +} + type runCtxKey struct{} func getRunCtx(ctx context.Context) *runContext { diff --git a/adk/runctx_test.go b/adk/runctx_test.go index bef1f44eb..7dbc7e32b 100644 --- a/adk/runctx_test.go +++ b/adk/runctx_test.go @@ -21,10 +21,12 @@ import ( "context" "encoding/gob" "errors" + "sync" "testing" "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/cloudwego/eino/schema" ) @@ -632,3 +634,215 @@ func TestGobEncodeStreamErrors(t *testing.T) { assert.NoError(t, err, "encoding runSession with WillRetryError stream should succeed") }) } + +func TestSanitizeRunContextForSessionCheckpointStripsSessionEvents(t *testing.T) { + output := &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage("kept", nil), + Role: schema.Assistant, + }, + } + kept := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + AgentName: "agent", + RunPath: []RunStep{{agentName: "root"}}, + Output: output, + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: "event-output", + Kind: SessionEventMessage, + Message: schema.AssistantMessage("kept", nil), + }, + }, + }, + TS: 10, + } + dropped := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: "event-session-only", + Kind: SessionEventSessionStatusRunning, + }, + }, + }, + TS: 11, + } + interrupt := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Action: &AgentAction{Interrupted: &InterruptInfo{Data: "pause"}}, + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: "event-interrupt", + Kind: SessionEventInterrupt, + }, + }, + }, + TS: 12, + } + session := newRunSession() + session.Values["k"] = "v" + session.Events = []*agentEventWrapper{kept, dropped, interrupt} + rc := &runContext{ + RootInput: &AgentInput{Messages: []*schema.Message{schema.UserMessage("q")}}, + RunPath: []RunStep{{agentName: "root"}}, + Session: session, + } + + sanitized := sanitizeRunContextForSessionCheckpoint[*schema.Message](rc) + + require.NotNil(t, sanitized) + require.NotSame(t, rc, sanitized) + require.NotSame(t, session, sanitized.Session) + require.Len(t, sanitized.Session.Events, 2) + assert.Nil(t, sanitized.Session.Events[0].SessionEventVariant) + assert.Same(t, output, sanitized.Session.Events[0].Output) + assert.NotNil(t, sanitized.Session.Events[1].Action.Interrupted) + assert.Nil(t, sanitized.Session.Events[1].SessionEventVariant) + assert.Equal(t, map[string]any{"k": "v"}, sanitized.Session.Values) + + assert.NotNil(t, kept.SessionEventVariant.Event, "sanitizer must not mutate the original output event") + assert.NotNil(t, dropped.SessionEventVariant.Event, "sanitizer must not mutate the original timeline event") + assert.NotNil(t, interrupt.SessionEventVariant.Event, "sanitizer must not mutate the original interrupt event") +} + +func TestSanitizeRunContextForSessionCheckpointTypedEvents(t *testing.T) { + output := &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: schema.UserAgenticMessage("kept"), + AgenticRole: schema.AgenticRoleTypeUser, + }, + } + events := []*typedAgentEventWrapper[*schema.AgenticMessage]{ + { + event: &TypedAgentEvent[*schema.AgenticMessage]{ + Output: output, + SessionEventVariant: &SessionEventVariant[*schema.AgenticMessage]{ + Event: &SessionEvent[*schema.AgenticMessage]{ + EventID: "typed-output", + Kind: SessionEventMessage, + Message: schema.UserAgenticMessage("kept"), + }, + }, + }, + TS: 20, + }, + { + event: &TypedAgentEvent[*schema.AgenticMessage]{ + SessionEventVariant: &SessionEventVariant[*schema.AgenticMessage]{ + Event: &SessionEvent[*schema.AgenticMessage]{ + EventID: "typed-session-only", + Kind: SessionEventSessionStatusRunning, + }, + }, + }, + TS: 21, + }, + } + session := newRunSession() + session.TypedEvents = &events + rc := &runContext{Session: session} + + sanitized := sanitizeRunContextForSessionCheckpoint[*schema.AgenticMessage](rc) + + store, ok := sanitized.Session.TypedEvents.(*[]*typedAgentEventWrapper[*schema.AgenticMessage]) + require.True(t, ok) + require.Len(t, *store, 1) + assert.Nil(t, (*store)[0].event.SessionEventVariant) + assert.Same(t, output, (*store)[0].event.Output) + assert.NotNil(t, events[0].event.SessionEventVariant.Event, "sanitizer must not mutate the original typed event") + assert.NotNil(t, events[1].event.SessionEventVariant.Event, "sanitizer must not mutate the original typed timeline event") +} + +func TestSanitizeRunContextForSessionCheckpointReducesEncodedPayload(t *testing.T) { + sessionEvent := &SessionEvent[*schema.Message]{ + EventID: "large-session-event", + Kind: SessionEventMessage, + Message: schema.AssistantMessage("large duplicated durable session payload", nil), + } + rc := &runContext{Session: newRunSession()} + rc.Session.Events = []*agentEventWrapper{ + { + AgentEvent: &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{Event: sessionEvent}, + }, + }, + { + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage("kept output", nil), + Role: schema.Assistant, + }, + }, + SessionEventVariant: &SessionEventVariant[*schema.Message]{Event: sessionEvent}, + }, + }, + } + + unsanitized, err := encodeRunnerCheckPointWithRunCtx(false, rc, nil, nil) + require.NoError(t, err) + sanitized, err := encodeRunnerCheckPointWithRunCtx( + false, + sanitizeRunContextForSessionCheckpoint[*schema.Message](rc), + nil, + nil, + ) + require.NoError(t, err) + assert.Less(t, len(sanitized), len(unsanitized)) + + _, decoded, _, err := runnerLoadCheckPointBytes(context.Background(), sanitized) + require.NoError(t, err) + require.Len(t, decoded.Session.Events, 1) + assert.Nil(t, decoded.Session.Events[0].SessionEventVariant) + assert.NotNil(t, decoded.Session.Events[0].Output) +} + +func TestSanitizeRunContextForSessionCheckpointPreservesLaneChain(t *testing.T) { + parentTimelineOnly := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{Event: &SessionEvent[*schema.Message]{Kind: SessionEventSessionStatusRunning}}, + }, + } + parentOutput := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage("parent", nil)}}, + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: "parent-output", + Kind: SessionEventMessage, + }, + }, + }, + } + childTimelineOnly := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{Event: &SessionEvent[*schema.Message]{Kind: SessionEventSessionStatusIdle}}, + }, + } + childOutput := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage("child", nil)}}, + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: "child-output", + Kind: SessionEventMessage, + }, + }, + }, + } + parent := &laneEvents{Events: []*agentEventWrapper{parentTimelineOnly, parentOutput}} + child := &laneEvents{Events: []*agentEventWrapper{childTimelineOnly, childOutput}, Parent: parent} + rc := &runContext{Session: &runSession{LaneEvents: child, valuesMtx: &sync.Mutex{}}} + + sanitized := sanitizeRunContextForSessionCheckpoint[*schema.Message](rc) + + require.NotNil(t, sanitized.Session.LaneEvents) + require.NotNil(t, sanitized.Session.LaneEvents.Parent) + require.Len(t, sanitized.Session.LaneEvents.Events, 1) + require.Len(t, sanitized.Session.LaneEvents.Parent.Events, 1) + assert.Nil(t, sanitized.Session.LaneEvents.Events[0].SessionEventVariant) + assert.Nil(t, sanitized.Session.LaneEvents.Parent.Events[0].SessionEventVariant) + assert.NotNil(t, childTimelineOnly.SessionEventVariant.Event, "sanitizer must not mutate original child lane") + assert.NotNil(t, parentTimelineOnly.SessionEventVariant.Event, "sanitizer must not mutate original parent lane") +} diff --git a/adk/runner.go b/adk/runner.go index a7d722e6f..f5bdd477c 100644 --- a/adk/runner.go +++ b/adk/runner.go @@ -17,11 +17,15 @@ package adk import ( + "bytes" "context" + "encoding/gob" "errors" "fmt" + "reflect" "runtime/debug" "sync" + "time" "github.com/cloudwego/eino/internal/core" "github.com/cloudwego/eino/internal/safe" @@ -56,6 +60,9 @@ type TypedRunner[M MessageType] struct { a TypedAgent[M] enableStreaming bool store CheckPointStore + sessionID string + sessionStore SessionEventStore[M] + sessionConfig *SessionConfig[M] } // Runner is the default runner type using *schema.Message. @@ -70,6 +77,10 @@ type TypedRunnerConfig[M MessageType] struct { EnableStreaming bool CheckPointStore CheckPointStore + + SessionID string + SessionStore SessionEventStore[M] + SessionConfig *SessionConfig[M] } // RunnerConfig is the default runner config type using *schema.Message. @@ -96,12 +107,15 @@ func NewTypedRunner[M MessageType](conf TypedRunnerConfig[M]) *TypedRunner[M] { enableStreaming: conf.EnableStreaming, a: conf.Agent, store: conf.CheckPointStore, + sessionID: conf.SessionID, + sessionStore: conf.SessionStore, + sessionConfig: conf.SessionConfig, } } func (r *TypedRunner[M]) Run(ctx context.Context, messages []M, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { - return typedRunnerRunImpl(r.a, r.enableStreaming, r.store, ctx, messages, opts...) + return typedRunnerRunImpl(r.a, r.enableStreaming, r.store, r.sessionID, r.sessionStore, r.sessionConfig, ctx, messages, opts...) } // Query is a convenience method that starts a new execution with a single user query string. @@ -150,17 +164,419 @@ func (r *TypedRunner[M]) ResumeWithParams(ctx context.Context, checkPointID stri func (r *TypedRunner[M]) resumeInternal(ctx context.Context, checkPointID string, resumeData map[string]any, opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) { - return typedRunnerResumeInternalImpl(r.a, r.store, ctx, checkPointID, resumeData, opts...) + return typedRunnerResumeInternalImpl(r.a, r.store, r.sessionID, r.sessionStore, r.sessionConfig, ctx, checkPointID, resumeData, opts...) +} + +type runnerSessionRunState[M MessageType] struct { + enabled bool + sessionID string + checkPointID *string + latestState *reconstructedSessionState[M] + sessionConfig SessionConfig[M] + sessionStore SessionEventStore[M] + sessionHandle sessionHandle[M] + checkPointStore CheckPointStore + initialTimeline []*SessionEvent[M] + // inputMessages are the caller-provided messages for this turn (before history prepend). + // Captured so the Runner can persist them as session events at turn start. + inputMessages []M +} + +func valueOrEmpty(v *string) string { + if v == nil { + return "" + } + return *v +} + +func isNilCheckPointStore(store CheckPointStore) bool { + if store == nil { + return true + } + v := reflect.ValueOf(store) + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice: + return v.IsNil() + default: + return false + } } -func typedRunnerRunImpl[M MessageType](a TypedAgent[M], enableStreaming bool, store CheckPointStore, ctx context.Context, messages []M, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { +func openRunnerSession[M MessageType]( + ctx context.Context, + store SessionEventStore[M], + sessionID string, + cfg SessionConfig[M], +) (*openSessionResult[M], error) { + if store == nil { + return nil, errors.New("adk: session store is nil") + } + deadline := timeNow().Add(cfg.SessionAcquireTimeout) + var lastErr error + for { + result, err := openLocalSession(ctx, store, &openSessionRequest{ + sessionID: sessionID, + }) + if err == nil { + if result == nil || result.handle == nil { + return nil, ErrSessionBusy + } + return result, nil + } + if !errors.Is(err, ErrSessionBusy) { + return nil, err + } + lastErr = err + if !timeNow().Before(deadline) { + return nil, lastErr + } + wait := 10 * time.Millisecond + var busy *SessionBusyError + if errors.As(err, &busy) && busy.ExpiresAt.After(timeNow()) { + until := time.Until(busy.ExpiresAt) + if until < wait { + wait = until + } + } + select { + case <-time.After(wait): + case <-ctx.Done(): + return nil, ctx.Err() + } + } +} + +func timeNow() time.Time { + return time.Now() +} + +func prepareRunnerSessionRun[M MessageType]( //nolint:revive // argument-limit + ctx context.Context, + checkPointStore CheckPointStore, + requestedCheckPointID *string, + sessionID string, + sessionStore SessionEventStore[M], + sessionConfig *SessionConfig[M], +) (*runnerSessionRunState[M], error) { + state := &runnerSessionRunState[M]{} + if isNilCheckPointStore(checkPointStore) { + checkPointStore = nil + } + if sessionID == "" || sessionStore == nil { + return state, nil + } + state.enabled = true + state.sessionID = sessionID + state.sessionStore = sessionStore + state.checkPointStore = checkPointStore + state.sessionConfig = normalizeSessionConfig(sessionConfig) + state.latestState = &reconstructedSessionState[M]{} + openResult, err := openRunnerSession[M](ctx, sessionStore, sessionID, state.sessionConfig) + if err != nil { + return nil, err + } + state.sessionHandle = openResult.handle + + reconstructResult, err := reconstructSessionState[M](ctx, state.sessionHandle, sessionID, defaultLoadPageSize) + if err != nil { + _ = state.sessionHandle.close(ctx) + return nil, fmt.Errorf("failed to reconstruct session[%s]: %w", sessionID, err) + } + if reconstructResult != nil && reconstructResult.state != nil { + state.latestState = reconstructResult.state + } + runningEvent := &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventSessionStatusRunning, + Lifecycle: &LifecycleEvent{State: SessionRunStateRunning}, + } + err = assignSessionEventID(ctx, runningEvent, state.sessionConfig.EventIDGenerator) + if err != nil { + _ = state.sessionHandle.close(ctx) + return nil, err + } + err = appendRunnerSessionControlEvent(ctx, state, runningEvent) + if err != nil { + _ = state.sessionHandle.close(ctx) + return nil, err + } + state.initialTimeline = append(state.initialTimeline, runningEvent) + + if isNilCheckPointStore(checkPointStore) { + return state, nil + } + checkPointID := sessionRunnerCheckpointID(sessionID) + if requestedCheckPointID != nil && *requestedCheckPointID != "" { + checkPointID = *requestedCheckPointID + } + state.checkPointID = &checkPointID + _, existed, err := loadRunnerSessionCheckpoint(ctx, checkPointStore, checkPointID) + if err != nil { + _ = state.sessionHandle.close(ctx) + return nil, err + } + if !existed { + return state, nil + } + // Pending checkpoint exists but caller chose Run (fresh turn) instead of Resume. + // We intentionally do NOT delete the checkpoint here — it remains available for a + // future Resume call. Session correctness is guaranteed by event log replay regardless + // of checkpoint presence. If this fresh turn completes successfully, finalize() will + // clean up the stale checkpoint at that point. + return state, nil +} + +func prepareRunnerSessionResume[M MessageType]( //nolint:revive // argument-limit + ctx context.Context, + checkPointStore CheckPointStore, + sessionID string, + sessionStore SessionEventStore[M], + sessionConfig *SessionConfig[M], + checkPointID string, +) (*runnerSessionRunState[M], string, error) { + state := &runnerSessionRunState[M]{} + if isNilCheckPointStore(checkPointStore) { + checkPointStore = nil + } + // Non-session-mode resume: explicit checkpoint ID, no session boot needed. + if checkPointID != "" && (sessionID == "" || sessionStore == nil) { + return state, checkPointID, nil + } + // Implicit session-mode resume requires both sessionID and sessionStore. + if checkPointID == "" && (sessionID == "" || sessionStore == nil) { + return nil, "", errors.New("failed to resume: checkpoint ID is empty") + } + state.enabled = true + state.sessionID = sessionID + state.sessionStore = sessionStore + state.checkPointStore = checkPointStore + state.sessionConfig = normalizeSessionConfig(sessionConfig) + state.latestState = &reconstructedSessionState[M]{} + openResult, err := openRunnerSession[M](ctx, sessionStore, sessionID, state.sessionConfig) + if err != nil { + return nil, "", err + } + state.sessionHandle = openResult.handle + + reconstructResult, err := reconstructSessionState[M](ctx, state.sessionHandle, sessionID, defaultLoadPageSize) + if err != nil { + _ = state.sessionHandle.close(ctx) + return nil, "", fmt.Errorf("failed to reconstruct session[%s]: %w", sessionID, err) + } + if reconstructResult != nil && reconstructResult.state != nil { + state.latestState = reconstructResult.state + } + // Pick the checkpoint ID: caller-provided takes precedence over the implicit + // session-scoped one. The session-scoped key still drives existence checks + // when the caller did not supply a checkpoint. + effectiveCheckPointID := checkPointID + if effectiveCheckPointID == "" { + effectiveCheckPointID = sessionRunnerCheckpointID(sessionID) + } + state.checkPointID = &effectiveCheckPointID + + // Existence check is only required for implicit session resume — the caller + // passing an explicit checkpoint ID has asserted the checkpoint should exist + // and any error will surface from the subsequent load. For implicit resume, + // the absence of a pending checkpoint is fatal and reported here. + _, existed, err := loadRunnerSessionCheckpoint(ctx, checkPointStore, effectiveCheckPointID) + if err != nil { + _ = state.sessionHandle.close(ctx) + return nil, "", err + } + if !existed { + _ = state.sessionHandle.close(ctx) + if checkPointID == "" { + return nil, "", fmt.Errorf("no pending session checkpoint for session %q", sessionID) + } + return nil, "", fmt.Errorf("checkpoint[%s] not exist", effectiveCheckPointID) + } + resumeEvent := &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventKind(SessionEventExtensionPrefix + "resume.request_started"), + Extension: &SessionExtensionEvent{}, + } + if err := assignSessionEventID(ctx, resumeEvent, state.sessionConfig.EventIDGenerator); err != nil { + _ = state.sessionHandle.close(ctx) + return nil, "", err + } + if err := appendRunnerSessionControlEvent(ctx, state, resumeEvent); err != nil { + _ = state.sessionHandle.close(ctx) + return nil, "", err + } + state.initialTimeline = append(state.initialTimeline, resumeEvent) + return state, effectiveCheckPointID, nil +} + +func appendRunnerSessionControlEvent[M MessageType]( + ctx context.Context, + state *runnerSessionRunState[M], + event *SessionEvent[M], +) error { + if state == nil || !state.enabled || state.sessionHandle == nil || event == nil { + return nil + } + if err := ValidateEmittedSessionEventKind(event); err != nil { + return err + } + err := state.sessionHandle.appendEvents(ctx, []*SessionEvent[M]{event}) + return err +} + +func appendRunnerSessionInputEvents[M MessageType]( + ctx context.Context, + state *runnerSessionRunState[M], + messages []M, +) error { + if state == nil || !state.enabled || state.sessionHandle == nil || len(messages) == 0 { + return nil + } + for _, msg := range messages { + se := makeInputSessionEvent[M](msg) + if err := assignSessionEventID(ctx, se, state.sessionConfig.EventIDGenerator); err != nil { + return err + } + if err := ValidateEmittedSessionEventKind(se); err != nil { + return err + } + if err := state.sessionHandle.appendEvents(ctx, []*SessionEvent[M]{se}); err != nil { + return err + } + state.initialTimeline = append(state.initialTimeline, se) + } + return nil +} + +func loadRunnerSessionCheckpoint(ctx context.Context, store CheckPointStore, checkPointID string) (*runnerSessionCheckpoint, bool, error) { + data, existed, err := store.Get(ctx, checkPointID) + if err != nil { + return nil, false, fmt.Errorf("failed to load session checkpoint[%s]: %w", checkPointID, err) + } + if !existed { + return nil, false, nil + } + cp, err := decodeRunnerSessionCheckpoint(data) + if err != nil { + return nil, false, fmt.Errorf("failed to decode session checkpoint[%s]: %w", checkPointID, err) + } + return cp, true, nil +} + +func runnerLoadCheckPointForSession(store CheckPointStore, ctx context.Context, checkPointID string, sessionMode bool) ( + context.Context, *runContext, *ResumeInfo, error) { + if !sessionMode { + return runnerLoadCheckPointImpl(store, ctx, checkPointID) + } + cp, existed, err := loadRunnerSessionCheckpoint(ctx, store, checkPointID) + if err != nil { + return nil, nil, nil, err + } + if !existed { + return nil, nil, nil, fmt.Errorf("checkpoint[%s] not exist", checkPointID) + } + return runnerLoadCheckPointBytes(ctx, cp.Payload) +} + +func runnerLoadCheckPointBytes(ctx context.Context, data []byte) ( + context.Context, *runContext, *ResumeInfo, error) { + data = preprocessADKCheckpoint(data) + s := &serialization{} + err := gob.NewDecoder(bytes.NewReader(data)).Decode(s) + if err != nil { + return nil, nil, nil, fmt.Errorf("failed to decode checkpoint: %w", err) + } + ctx = core.PopulateInterruptState(ctx, s.InterruptID2Address, s.InterruptID2State) + return ctx, s.RunCtx, &ResumeInfo{ + EnableStreaming: s.EnableStreaming, + InterruptInfo: s.Info, + }, nil +} + +func deleteCheckPointIfSupported(ctx context.Context, store CheckPointStore, checkPointID string) error { + if isNilCheckPointStore(store) { + return nil + } + if deleter, ok := store.(CheckPointDeleter); ok { + return deleter.Delete(ctx, checkPointID) + } + return nil +} + +func saveRunnerCheckpoint[M MessageType]( //nolint:revive // argument-limit + enableStreaming bool, + store CheckPointStore, + ctx context.Context, + checkPointID string, + info *InterruptInfo, + is *core.InterruptSignal, + sessionState *runnerSessionRunState[M], +) error { + if sessionState == nil || !sessionState.enabled { + return runnerSaveCheckPointImpl(enableStreaming, store, ctx, checkPointID, info, is) + } + if isNilCheckPointStore(store) { + return nil + } + payload, err := encodeRunnerCheckPointWithRunCtx( + enableStreaming, + sanitizeRunContextForSessionCheckpoint[M](getRunCtx(ctx)), + info, + is, + ) + if err != nil { + return err + } + data, err := encodeRunnerSessionCheckpoint(&runnerSessionCheckpoint{ + SessionID: sessionState.sessionID, + CheckPointID: checkPointID, + Payload: payload, + }) + if err != nil { + return err + } + return store.Set(ctx, checkPointID, data) +} + +func typedRunnerRunImpl[M MessageType](a TypedAgent[M], enableStreaming bool, store CheckPointStore, sessionID string, sessionStore SessionEventStore[M], sessionConfig *SessionConfig[M], ctx context.Context, messages []M, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { //nolint:revive // argument-limit o := getCommonOptions(nil, opts...) + exposeTimelineEvents := o.enableTimelineEvents + + sessionState, err := prepareRunnerSessionRun[M](ctx, store, o.checkPointID, sessionID, sessionStore, sessionConfig) + if err != nil { + return errorIterator[M](err) + } + if sessionState.enabled { + // Capture caller-provided messages BEFORE prepending history. These will be + // emitted as session events at turn start so they appear in the event log. + sessionState.inputMessages = append([]M{}, messages...) + messages = append(append([]M{}, sessionState.latestState.Messages...), sessionState.inputMessages...) + // Assign IDs before messages can be both persisted and inspected by + // middleware, avoiding concurrent lazy ID mutation during event snapshotting. + for _, msg := range messages { + EnsureMessageID(msg) + } + if err := appendRunnerSessionInputEvents(ctx, sessionState, sessionState.inputMessages); err != nil { + _ = sessionState.sessionHandle.close(ctx) + return errorIterator[M](err) + } + sessionState.inputMessages = nil + opts = append(opts, withEnableSessionEvents()) + opts = append(opts, withEnableInternalTimelineEvents()) + opts = append(opts, withInitialModelContext(&ModelContextEvent{ + ToolInfos: sessionState.latestState.ToolInfos, + DeferredToolInfos: sessionState.latestState.DeferredToolInfos, + }, sessionState.latestState.sawModelContext)) + } input := &TypedAgentInput[M]{ Messages: messages, EnableStreaming: enableStreaming, } + if sessionState.enabled { + ctx = contextWithSessionEventIDGenerator[M](ctx, sessionState.sessionConfig.EventIDGenerator) + } + var zero M if _, ok := any(zero).(*schema.Message); ok { concreteAgent, _ := any(a).(Agent) @@ -174,12 +590,19 @@ func typedRunnerRunImpl[M MessageType](a TypedAgent[M], enableStreaming bool, st iter := fa.Run(ctx, concreteInput, opts...) - if store == nil && o.cancelCtx == nil { + // Short-circuit: no checkpoint to save, no cancel to handle, and no need to + // strip session-internal fields (enableSessionEvents means the caller wants + // them). The intermediate iterator pair adds no value in this case. + if store == nil && o.cancelCtx == nil && exposeTimelineEvents && !sessionState.enabled { return any(iter).(*AsyncIterator[*TypedAgentEvent[M]]) } niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() - go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(iter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, o.checkPointID, o.cancelCtx) + checkPointID := o.checkPointID + if sessionState.checkPointID != nil { + checkPointID = sessionState.checkPointID + } + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(iter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, checkPointID, o.cancelCtx, exposeTimelineEvents, sessionState) return niter } @@ -193,23 +616,49 @@ func typedRunnerRunImpl[M MessageType](a TypedAgent[M], enableStreaming bool, st iter := fa.Run(ctx, input, opts...) - if store == nil && o.cancelCtx == nil { + // Short-circuit: no checkpoint to save, no cancel to handle, and no need to + // strip session-internal fields (enableSessionEvents means the caller wants + // them). The intermediate iterator pair adds no value in this case. + if store == nil && o.cancelCtx == nil && exposeTimelineEvents && !sessionState.enabled { return iter } niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() - go typedRunnerHandleIterImpl(enableStreaming, store, ctx, iter, gen, o.checkPointID, o.cancelCtx) + checkPointID := o.checkPointID + if sessionState.checkPointID != nil { + checkPointID = sessionState.checkPointID + } + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, iter, gen, checkPointID, o.cancelCtx, exposeTimelineEvents, sessionState) return niter } -func typedRunnerResumeInternalImpl[M MessageType](a TypedAgent[M], store CheckPointStore, ctx context.Context, checkPointID string, resumeData map[string]any, //nolint:revive // argument-limit +func typedRunnerResumeInternalImpl[M MessageType](a TypedAgent[M], store CheckPointStore, sessionID string, sessionStore SessionEventStore[M], sessionConfig *SessionConfig[M], ctx context.Context, checkPointID string, resumeData map[string]any, //nolint:revive // argument-limit opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) { - if store == nil { + if isNilCheckPointStore(store) { return nil, fmt.Errorf("failed to resume: store is nil") } - ctx, runCtx, resumeInfo, err := runnerLoadCheckPointImpl(store, ctx, checkPointID) + o := getCommonOptions(nil, opts...) + exposeTimelineEvents := o.enableTimelineEvents + sessionState, effectiveCheckPointID, err := prepareRunnerSessionResume[M](ctx, store, sessionID, sessionStore, sessionConfig, checkPointID) + if err != nil { + return nil, err + } + checkPointID = effectiveCheckPointID + if sessionState.enabled { + opts = append(opts, withEnableSessionEvents()) + opts = append(opts, withEnableInternalTimelineEvents()) + opts = append(opts, withInitialModelContext(&ModelContextEvent{ + ToolInfos: sessionState.latestState.ToolInfos, + DeferredToolInfos: sessionState.latestState.DeferredToolInfos, + }, sessionState.latestState.sawModelContext)) + } + + ctx, runCtx, resumeInfo, err := runnerLoadCheckPointForSession(store, ctx, checkPointID, sessionState.enabled) if err != nil { + if sessionState != nil && sessionState.enabled && sessionState.sessionHandle != nil { + _ = sessionState.sessionHandle.close(ctx) + } return nil, fmt.Errorf("failed to load from checkpoint: %w", err) } @@ -219,7 +668,6 @@ func typedRunnerResumeInternalImpl[M MessageType](a TypedAgent[M], store CheckPo // running in, and any new checkpoint written during this resume must preserve it. enableStreaming := resumeInfo.EnableStreaming - o := getCommonOptions(nil, opts...) if o.sharedParentSession { parentSession := getSession(ctx) if parentSession != nil { @@ -237,6 +685,10 @@ func typedRunnerResumeInternalImpl[M MessageType](a TypedAgent[M], store CheckPo ctx = setRunCtx(ctx, runCtx) AddSessionValues(ctx, o.sessionValues) + if sessionState.enabled { + ctx = contextWithSessionEventIDGenerator[M](ctx, sessionState.sessionConfig.EventIDGenerator) + } + if len(resumeData) > 0 { ctx = core.BatchResumeWithData(ctx, resumeData) } @@ -245,31 +697,37 @@ func typedRunnerResumeInternalImpl[M MessageType](a TypedAgent[M], store CheckPo if _, ok := any(zero).(*schema.Message); ok { concreteAgent, _ := any(a).(Agent) fa := toFlowAgent(ctx, concreteAgent) - ra, ok := Agent(fa).(ResumableAgent) + ra, ok := any(fa).(ResumableAgent) if !ok { + if sessionState.enabled && sessionState.sessionHandle != nil { + _ = sessionState.sessionHandle.close(ctx) + } return nil, fmt.Errorf("agent %T does not support resume", a) } aIter := ra.Resume(ctx, resumeInfo, opts...) niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() - go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(aIter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, &checkPointID, o.cancelCtx) + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(aIter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, &checkPointID, o.cancelCtx, exposeTimelineEvents, sessionState) return niter, nil } fa := toTypedFlowAgent(a) - ra, ok := TypedAgent[M](fa).(TypedResumableAgent[M]) + ra, ok := any(fa).(TypedResumableAgent[M]) if !ok { + if sessionState.enabled && sessionState.sessionHandle != nil { + _ = sessionState.sessionHandle.close(ctx) + } return nil, fmt.Errorf("agent %T does not support resume", a) } aIter := ra.Resume(ctx, resumeInfo, opts...) niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() - go typedRunnerHandleIterImpl(enableStreaming, store, ctx, aIter, gen, &checkPointID, o.cancelCtx) + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, aIter, gen, &checkPointID, o.cancelCtx, exposeTimelineEvents, sessionState) return niter, nil } -func typedRunnerHandleIterImpl[M MessageType](enableStreaming bool, store CheckPointStore, ctx context.Context, aIter *AsyncIterator[*TypedAgentEvent[M]], //nolint:revive // argument-limit - gen *AsyncGenerator[*TypedAgentEvent[M]], checkPointID *string, cancelCtx *cancelContext) { +func typedRunnerHandleIterImpl[M MessageType](enableStreaming bool, store CheckPointStore, ctx context.Context, aIter *AsyncIterator[*TypedAgentEvent[M]], //nolint:revive,cyclop,funlen // argument-limit; event loop branches by event kind + gen *AsyncGenerator[*TypedAgentEvent[M]], checkPointID *string, cancelCtx *cancelContext, enableTimelineEvents bool, sessionState *runnerSessionRunState[M]) { defer func() { panicErr := recover() if panicErr != nil { @@ -280,31 +738,221 @@ func typedRunnerHandleIterImpl[M MessageType](enableStreaming bool, store CheckP gen.Close() }() var ( - interruptSignal *core.InterruptSignal - legacyData any + interruptSignal *core.InterruptSignal + interruptContexts []*InterruptCtx + legacyData any + interrupted bool + cancelled bool + retryExhausted bool + terminalErr error + persister *sessionEventPersister[M] + persistErr error + + // pendingCheckpoint defers checkpoint save to finalize() so the persister + // can flush enqueued events first. Writing the checkpoint before the flush + // completes risks a checkpoint that references events not yet durable. + pendingCheckpoint *deferredRunnerCheckpoint ) + if sessionState != nil && sessionState.enabled { + persister = newSessionEventPersister[M](ctx, sessionState.sessionHandle, sessionState.sessionID) + } + if enableTimelineEvents && sessionState != nil && sessionState.enabled { + for _, se := range sessionState.initialTimeline { + if se != nil { + gen.Send(&TypedAgentEvent[M]{SessionEventVariant: &SessionEventVariant[M]{SessionID: sessionState.sessionID, Event: se}}) + } + } + } + setPersistErr := func(err error) { + if err != nil && persistErr == nil { + persistErr = err + } + } + enqueueAsyncSessionEvent := func(se *SessionEvent[M]) error { + if persister == nil || se == nil { + return nil + } + if err := persister.enqueueAsync(se); err != nil { + setPersistErr(err) + return err + } + return nil + } + commitSessionBoundary := func(se *SessionEvent[M]) error { + if persister == nil || se == nil { + return nil + } + if err := persister.commitBoundary(se); err != nil { + setPersistErr(err) + return err + } + return nil + } + persistSessionEvent := func(se *SessionEvent[M]) error { + if persister == nil || se == nil { + return nil + } + if err := ValidateEmittedSessionEventKind(se); err != nil { + setPersistErr(err) + return err + } + if isSessionDurableBoundaryKind(se.Kind) { + return commitSessionBoundary(se) + } + return enqueueAsyncSessionEvent(se) + } + sendTimelineEvent := func(se *SessionEvent[M]) bool { + if se == nil { + return false + } + if se.EventID == "" { + if err := assignSessionEventIDFromContext(ctx, se); err != nil { + setPersistErr(err) + return false + } + } + if se.Timestamp.IsZero() { + se.Timestamp = newEventTimestamp() + } + if err := persistSessionEvent(se); err != nil { + return false + } + event := &TypedAgentEvent[M]{SessionEventVariant: &SessionEventVariant[M]{SessionID: sessionState.sessionID, Event: se}} + if enableTimelineEvents { + gen.Send(event) + } + return true + } + reserveMessageStreamRef := func(event *TypedAgentEvent[M]) (*MessageStreamRef, error) { + if event != nil && event.SessionEventVariant != nil && event.SessionEventVariant.MessageStreamRef != nil { + ref := event.SessionEventVariant.MessageStreamRef + if ref.Timestamp.IsZero() { + ref.Timestamp = newEventTimestamp() + } + ref.Kind = SessionEventMessage + if ref.EventID == "" { + draft := &SessionEvent[M]{ + Timestamp: ref.Timestamp, + Kind: SessionEventMessage, + } + if err := assignSessionEventID(ctx, draft, sessionState.sessionConfig.EventIDGenerator); err != nil { + setPersistErr(err) + return nil, err + } + ref.EventID = draft.EventID + ref.Timestamp = draft.Timestamp + } + return ref, nil + } + draft := &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventMessage, + } + if err := assignSessionEventID(ctx, draft, sessionState.sessionConfig.EventIDGenerator); err != nil { + setPersistErr(err) + return nil, err + } + return &MessageStreamRef{ + EventID: draft.EventID, + Timestamp: draft.Timestamp, + Kind: SessionEventMessage, + }, nil + } + toSessionEventCheckedWithGenerator := func(event *TypedAgentEvent[M]) (*SessionEvent[M], error) { + se, err := toSessionEventChecked(event) + if err == nil || event == nil || event.SessionEventVariant != nil || + event.Output == nil || event.Output.MessageOutput == nil || + isNilMessage(event.Output.MessageOutput.Message) { + return se, err + } + draft := &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventMessage, + Message: event.Output.MessageOutput.Message, + } + if idErr := assignSessionEventID(ctx, draft, sessionState.sessionConfig.EventIDGenerator); idErr != nil { + return nil, idErr + } + event.SessionEventVariant = &SessionEventVariant[M]{SessionID: sessionState.sessionID, Event: draft} + return draft, NormalizeSessionEventKind(draft) + } + // saveCheckpointNow is the path used when no session persister is active — + // the checkpoint is written immediately because there are no queued events + // to flush. In session mode, the same payload is captured into + // pendingCheckpoint and committed inside finalize() after persister.closeAndWait. + saveCheckpointNow := func(info *InterruptInfo, sig *core.InterruptSignal, errLabel string) { + if checkPointID == nil { + return + } + if info == nil { + info = &InterruptInfo{} + } + info.CheckPointID = *checkPointID + if persister != nil { + pendingCheckpoint = &deferredRunnerCheckpoint{info: info, signal: sig, errLabel: errLabel} + return + } + if err := saveRunnerCheckpoint(enableStreaming, store, ctx, *checkPointID, info, sig, sessionState); err != nil { + gen.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("%s: %w", errLabel, err)}) + } + } + for { event, ok := aIter.Next() if !ok { break } + fromOtherSession := event.SessionEventVariant != nil && + sessionState != nil && sessionState.enabled && + event.SessionEventVariant.SessionID != "" && + event.SessionEventVariant.SessionID != sessionState.sessionID + if !fromOtherSession && event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil { + gen := DefaultSessionEventIDGenerator[M] + if sessionState != nil && sessionState.enabled { + gen = sessionState.sessionConfig.EventIDGenerator + if gen == nil { + gen = DefaultSessionEventIDGenerator[M] + } + } + if _, err := normalizeAgentSessionEventWithAssigner(event, func(draft *SessionEvent[M]) (string, error) { + return gen(ctx, draft) + }); err != nil { + setPersistErr(err) + event.Err = err + } + } + if err := validateAgentSessionEventIdentity(event); err != nil { + setPersistErr(err) + event.Err = err + } if event.Err != nil { + var retryErr *RetryExhaustedError + if errors.As(event.Err, &retryErr) { + retryExhausted = true + } var cancelErr *CancelError if errors.As(event.Err, &cancelErr) { + cancelled = true if cancelCtx != nil && cancelCtx.isRoot() && cancelCtx.shouldCancel() { cancelCtx.markCancelHandled() } if cancelErr.interruptSignal != nil && checkPointID != nil { cancelErr.InterruptContexts = core.ToInterruptContexts(cancelErr.interruptSignal, allowedAddressSegmentTypes) - err := runnerSaveCheckPointImpl(enableStreaming, store, ctx, *checkPointID, &InterruptInfo{}, cancelErr.interruptSignal) - if err != nil { - gen.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("failed to save checkpoint on cancel: %w", err)}) + saveCheckpointNow(&InterruptInfo{}, cancelErr.interruptSignal, "failed to save checkpoint on cancel") + } + if !enableTimelineEvents { + event = stripSessionEventFields(event) + if event == nil { + break } } gen.Send(event) break } + if terminalErr == nil { + terminalErr = event.Err + } } if event.Action != nil && event.Action.internalInterrupted != nil { @@ -312,7 +960,7 @@ func typedRunnerHandleIterImpl[M MessageType](enableStreaming bool, store CheckP panic("multiple interrupt actions should not happen in Runner") } interruptSignal = event.Action.internalInterrupted - interruptContexts := core.ToInterruptContexts(interruptSignal, allowedAddressSegmentTypes) + interruptContexts = core.ToInterruptContexts(interruptSignal, allowedAddressSegmentTypes) event = &TypedAgentEvent[M]{ AgentName: event.AgentName, RunPath: event.RunPath, @@ -320,23 +968,273 @@ func typedRunnerHandleIterImpl[M MessageType](enableStreaming bool, store CheckP Action: &AgentAction{ Interrupted: &InterruptInfo{ Data: event.Action.Interrupted.Data, + CheckPointID: valueOrEmpty(checkPointID), InterruptContexts: interruptContexts, }, internalInterrupted: interruptSignal, }, } legacyData = event.Action.Interrupted.Data + interrupted = true if checkPointID != nil { - err := runnerSaveCheckPointImpl(enableStreaming, store, ctx, *checkPointID, &InterruptInfo{ - Data: legacyData, - }, interruptSignal) - if err != nil { - gen.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("failed to save checkpoint: %w", err)}) + saveCheckpointNow(&InterruptInfo{Data: legacyData}, interruptSignal, "failed to save checkpoint") + } + } + + liveDelivered := false + if persister != nil { + // Skip persistence (but not live delivery) for events owned by a + // different session (inner agent events forwarded via AgentTool). + if !fromOtherSession { + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming && event.Output.MessageOutput.MessageStream != nil { + ref, err := reserveMessageStreamRef(event) + if err != nil { + continue + } + // Streaming output is split into two stream copies: copies[1] is + // rewritten onto the live event and sent immediately so live + // consumers see no extra latency. The message boundary is committed + // after copies[0] is drained and fully materialized. + copies := event.Output.MessageOutput.MessageStream.Copy(2) + liveOutput := *event.Output + liveMV := *event.Output.MessageOutput + liveMV.MessageStream = copies[1] + + liveOutput.MessageOutput = &liveMV + event.Output = &liveOutput + event.SessionEventVariant = &SessionEventVariant[M]{SessionID: sessionState.sessionID, MessageStreamRef: ref} + liveEvent := event + if !enableTimelineEvents { + liveEvent = stripSessionEventFields(liveEvent) + } + if liveEvent != nil { + gen.Send(liveEvent) + } + liveDelivered = true + + persistedMsg, hasChunks, streamErr, err := materializeMessageStreamPrefix(copies[0]) + if err != nil { + // Prefix projection is best-effort replay data. A concat failure + // should not fail the turn after the source stream already failed. + continue + } + if streamErr != nil { + if hasChunks { + _ = persistSessionEvent(&SessionEvent[M]{ + EventID: ref.EventID, + Timestamp: ref.Timestamp, + Kind: SessionEventMessageStreamIncomplete, + MessageStreamIncomplete: &MessageStreamIncompleteEvent[M]{ + Message: persistedMsg, + Error: streamErr.Error(), + }, + }) + } + continue + } + if !hasChunks { + continue + } + + _ = persistSessionEvent(&SessionEvent[M]{ + EventID: ref.EventID, + Timestamp: ref.Timestamp, + Kind: SessionEventMessage, + Message: persistedMsg, + }) + } else { + // Non-streaming events go through toSessionEvent directly. + se, err := toSessionEventCheckedWithGenerator(event) + if err != nil { + setPersistErr(err) + se = nil + } + if se != nil { + if err := persistSessionEvent(se); err != nil { + continue + } + // Backfill SessionEventVariant onto the live event so downstream + // consumers (TurnLoop/onAgentEvents) see message events + // with their persisted SessionEvent identity, consistent + // with how span events are already delivered. + event.SessionEventVariant = &SessionEventVariant[M]{SessionID: sessionState.sessionID, Event: se} + } } } } + if liveDelivered { + continue + } + + if !enableTimelineEvents { + event = stripSessionEventFields(event) + if event == nil { + continue + } + } gen.Send(event) } + if persister != nil { + stopReason := "end_turn" + switch { + case interrupted: + stopReason = "interrupted" + case cancelled: + stopReason = "cancelled" + case retryExhausted: + stopReason = "retries_exhausted" + case persistErr != nil: + stopReason = "failed" + case terminalErr != nil: + stopReason = "failed" + } + if stopReason == "failed" { + errMsg := "" + if persistErr != nil { + errMsg = persistErr.Error() + } else if terminalErr != nil { + errMsg = terminalErr.Error() + } + sendTimelineEvent(&SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventSessionError, + Error: &SessionErrorEvent{Type: SessionErrorTypeFatal, Message: errMsg}, + }) + } + if interrupted { + sendTimelineEvent(&SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventInterrupt, + Interrupt: buildInterruptEvent(interruptContexts), + }) + } + if cancelled { + sendTimelineEvent(&SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventCancel, + Cancel: &CancelEvent{Reason: "cancelled"}, + }) + } + sendTimelineEvent(&SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventSessionStatusIdle, + Lifecycle: &LifecycleEvent{State: SessionRunStateIdle, StopReason: &StopReason{Type: stopReason}}, + }) + res := &sessionTurnResult[M]{ + persister: persister, + persistErr: persistErr, + interrupted: interrupted, + cancelled: cancelled, + terminalErr: terminalErr, + sessionState: sessionState, + store: store, + checkPointID: checkPointID, + enableStreaming: enableStreaming, + pendingCheckpoint: pendingCheckpoint, + } + if err := res.finalize(ctx); err != nil { + gen.Send(&TypedAgentEvent[M]{Err: err}) + } + } +} + +func buildInterruptEvent( + contexts []*InterruptCtx, +) *InterruptEvent { + event := &InterruptEvent{ + Contexts: make([]*InterruptContext, 0, len(contexts)), + } + for _, ctx := range contexts { + if ctx == nil { + continue + } + aic := &InterruptContext{ + InterruptID: ctx.ID, + Info: ctx.Info, + } + if toolUseID := extractToolUseID(ctx); toolUseID != "" { + aic.ToolUseID = toolUseID + } + event.Contexts = append(event.Contexts, aic) + } + return event +} + +func extractToolUseID(ctx *InterruptCtx) string { + for _, segment := range ctx.Address { + if segment.Type != AddressSegmentTool { + continue + } + if segment.SubID != "" { + return segment.SubID + } + return segment.ID + } + return "" +} + +// deferredRunnerCheckpoint captures the arguments needed to persist a runner +// checkpoint after the session event persister has flushed. Saving the +// checkpoint earlier would risk a checkpoint that references events not yet +// durable in the SessionEventStore. +type deferredRunnerCheckpoint struct { + info *InterruptInfo + signal *core.InterruptSignal + errLabel string +} + +// sessionTurnResult bundles the accumulated state from a Runner turn's event +// loop and drives the session commit-or-abort decision. +type sessionTurnResult[M MessageType] struct { + persister *sessionEventPersister[M] + persistErr error + interrupted bool + cancelled bool + terminalErr error + sessionState *runnerSessionRunState[M] + store CheckPointStore + checkPointID *string + enableStreaming bool + pendingCheckpoint *deferredRunnerCheckpoint +} + +func (r *sessionTurnResult[M]) finalize(ctx context.Context) error { + defer func() { + if r.sessionState != nil && r.sessionState.sessionHandle != nil { + _ = r.sessionState.sessionHandle.close(ctx) + } + }() + if err := r.persister.closeAndWait(); err != nil && r.persistErr == nil { + r.persistErr = err + } + // For interrupt/cancel paths, the checkpoint write is deferred until here so + // the persister's queued events are durable BEFORE the checkpoint references + // them. If event persistence failed, skip the checkpoint write entirely so + // resume cannot load a checkpoint that points to a corrupt event log. + if r.pendingCheckpoint != nil && r.checkPointID != nil { + if r.persistErr != nil { + return fmt.Errorf("%s: skipped because session event persistence failed: %w", r.pendingCheckpoint.errLabel, r.persistErr) + } + if err := saveRunnerCheckpoint(r.enableStreaming, r.store, ctx, *r.checkPointID, r.pendingCheckpoint.info, r.pendingCheckpoint.signal, r.sessionState); err != nil { + return fmt.Errorf("%s: %w", r.pendingCheckpoint.errLabel, err) + } + } + if r.persistErr != nil { + return fmt.Errorf("failed to persist session events: %w", r.persistErr) + } + if r.interrupted || r.cancelled { + return nil + } + if r.terminalErr != nil { + return nil + } + if r.checkPointID != nil && !isNilCheckPointStore(r.store) { + if err := deleteCheckPointIfSupported(ctx, r.store, *r.checkPointID); err != nil { + return fmt.Errorf("failed to delete session checkpoint: %w", err) + } + } + return nil } diff --git a/adk/session.go b/adk/session.go new file mode 100644 index 000000000..773f578bd --- /dev/null +++ b/adk/session.go @@ -0,0 +1,1673 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "fmt" + "strings" + "sync" + "time" + + "github.com/google/uuid" + + "github.com/cloudwego/eino/schema" +) + +const ( + defaultLoadPageSize = 100 + defaultSessionAcquireTimeout = 5 * time.Second +) + +// ErrInvalidEventID is returned by AppendEvents when a SessionEvent has an +// empty EventID. Protocol-level: persisters MUST NOT retry. +// +// Services accept any non-empty string as EventID. UUIDv4 is the Runner-side +// allocation format, but service implementations treat EventID as opaque. +var ErrInvalidEventID = errors.New("adk: session event has invalid event_id") + +// ErrSessionEventIDGeneratorEmpty is returned by assignSessionEventID when the +// configured SessionEventIDGenerator returns an empty event_id. This is a +// generator-side contract violation surfaced before AppendEvents is called. +// It is a separate sentinel from ErrInvalidEventID (store-side) so callers +// can distinguish "application generator violated its contract" from "store +// rejected an event_id". +var ErrSessionEventIDGeneratorEmpty = errors.New("adk: session event id generator returned empty event id") + +// ErrEventIDOutOfRange is returned by LoadEvents when +// LoadSessionEventsRequest.After references an event_id that does not exist in +// the session log. Callers can detect this and fall back to a full reload. +var ErrEventIDOutOfRange = errors.New("adk: session event id out of range") + +var ErrRollbackTargetNotFound = errors.New("adk: rollback target event not found") +var ErrInvalidRollbackTarget = errors.New("adk: invalid rollback target") +var ErrRollbackTargetInactive = errors.New("adk: rollback target is not active") +var ErrSessionHeadChanged = errors.New("adk: session committed idle head changed") +var ErrSessionBusy = errors.New("adk: session already has an active handle") +var ErrDuplicateEventID = errors.New("adk: duplicate session event_id") + +type SessionBusyError struct { + ExpiresAt time.Time +} + +func (e *SessionBusyError) Error() string { return ErrSessionBusy.Error() } +func (e *SessionBusyError) Unwrap() error { return ErrSessionBusy } + +const ( + sessionRunnerCheckpointSuffix = "/runner_checkpoint" +) + +// SessionEventStore is the provider-facing interface for a typed append-only +// session event log. Runner coordinates process-local single-writer access for +// a session before calling AppendEvents. +type SessionEventStore[M MessageType] interface { + LoadEvents(ctx context.Context, sessionID string, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[M], error) + AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[M]) error +} + +type openSessionRequest struct { + sessionID string +} + +type openSessionResult[M MessageType] struct { + handle sessionHandle[M] +} + +type sessionHandle[M MessageType] interface { + loadEvents(ctx context.Context, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[M], error) + appendEvents(ctx context.Context, events []*SessionEvent[M]) error + close(ctx context.Context) error +} + +// LoadSessionEventsRequest configures typed event loading pagination and direction. +type LoadSessionEventsRequest struct { + // After is the last-seen event_id used as an exclusive append-position cursor. + After string + // Limit is the maximum number of events to return. 0 means no limit. + Limit int + // Reverse, when true, returns events in newest-first order. + Reverse bool + // Kinds filters events by their Kind field. Empty means no kind filter. + Kinds []SessionEventKind +} + +// LoadSessionEventsResult is the response from SessionEventStore.LoadEvents. +type LoadSessionEventsResult[M MessageType] struct { + // Events are typed SessionEvent values owned by the caller. + Events []*SessionEvent[M] + // Next is the event_id of the last event in this page in the direction of travel. + Next string +} + +// SessionEvent is the JSON-serializable persistence format for session events. +// SessionEvent is session-local; stores receive the owning session id as a +// first-class argument. The pair (session_id, event_id) globally identifies a +// persisted event. +// Exactly one semantic content field is active per event. The MessagesReplaced field +// uses pointer-to-slice semantics (nil = absent, non-nil = active replacement). +type SessionEvent[M MessageType] struct { + // EventID is the canonical, session-unique identity of this event. + // Assigned exactly once by the Runner at event materialization + // (in makeInputSessionEvent / toSessionEvent). Persister-level retries + // re-send the same payload bytes and therefore the same EventID, which is + // what enables AppendEvents idempotency. Runner-allocated EventIDs are + // UUIDv4 strings; SessionEventStore implementations treat EventID as an opaque + // non-empty string and do NOT enforce UUIDv4 format. + // + // Distinct from MessageUpdatedEvent.MessageID: EventID identifies the + // session event envelope; MessageID identifies a logical message inside + // the session message array. + EventID string `json:"event_id"` + + // Timestamp is inherited from the source AgentEvent and represents the event + // occurrence time, not the SessionEventStore persistence time. + Timestamp time.Time `json:"timestamp,omitempty"` + + Kind SessionEventKind `json:"kind,omitempty"` + + Message M `json:"message,omitempty"` + MessageStreamIncomplete *MessageStreamIncompleteEvent[M] `json:"message_stream_incomplete,omitempty"` + MessagesReplaced *[]M `json:"messages_replaced,omitempty"` + MessageUpdated *MessageUpdatedEvent[M] `json:"message_updated,omitempty"` + MessageInserted *MessageInsertedEvent[M] `json:"message_inserted,omitempty"` + MessagesDeleted *MessagesDeletedEvent `json:"messages_deleted,omitempty"` + ModelContext *ModelContextEvent `json:"model_context,omitempty"` + Rollback *SessionRollbackEvent `json:"rollback,omitempty"` + + Lifecycle *LifecycleEvent `json:"lifecycle,omitempty"` + Error *SessionErrorEvent `json:"error,omitempty"` + Span *SpanEvent `json:"span,omitempty"` + + Cancel *CancelEvent `json:"cancel,omitempty"` + Interrupt *InterruptEvent `json:"interrupt,omitempty"` + Extension *SessionExtensionEvent `json:"extension,omitempty"` +} + +// SessionEventVariant is the live AgentEvent envelope for session-related +// metadata. SessionID is live ownership metadata only; it is intentionally not +// part of durable SessionEvent payloads. +// +// Invariant: exactly one of Event or MessageStreamRef must be set. +// Event carries a fully materialized SessionEvent; MessageStreamRef carries +// only the reserved durable identity for a streaming message whose content +// remains in Output.MessageOutput.MessageStream. +type SessionEventVariant[M MessageType] struct { + SessionID string + + Event *SessionEvent[M] + MessageStreamRef *MessageStreamRef +} + +// MessageStreamRef carries the durable identity metadata for a streaming +// message whose content remains in Output.MessageOutput.MessageStream. It is +// carried by SessionEventVariant for live events. The runner later reuses this +// identity when it drains its persistence copy of the stream and writes the +// resulting message as a SessionEvent. +type MessageStreamRef struct { + EventID string + Timestamp time.Time + Kind SessionEventKind +} + +type SessionEventKind string + +const ( + SessionEventMessage SessionEventKind = "message" + SessionEventMessageStreamIncomplete SessionEventKind = "message_stream_incomplete" + SessionEventMessagesReplaced SessionEventKind = "messages_replaced" + SessionEventMessageUpdated SessionEventKind = "message_updated" + SessionEventMessageInserted SessionEventKind = "message_inserted" + SessionEventMessagesDeleted SessionEventKind = "messages_deleted" + SessionEventModelContext SessionEventKind = "model_context" + SessionEventRollback SessionEventKind = "rollback" + + SessionEventSessionStatusRunning SessionEventKind = "session.status_running" + SessionEventSessionStatusIdle SessionEventKind = "session.status_idle" + SessionEventSessionError SessionEventKind = "session.error" + + SessionEventSpanModelRequestStart SessionEventKind = "span.model_request_start" + SessionEventSpanModelRequestEnd SessionEventKind = "span.model_request_end" + SessionEventSpanToolCallStart SessionEventKind = "span.tool_call_start" + SessionEventSpanToolCallEnd SessionEventKind = "span.tool_call_end" + + SessionEventCancel SessionEventKind = "cancel" + SessionEventInterrupt SessionEventKind = "interrupt" + + SessionEventExtensionPrefix = "x." +) + +var knownSessionEventKinds = map[SessionEventKind]struct{}{ + SessionEventMessage: {}, + SessionEventMessageStreamIncomplete: {}, + SessionEventMessagesReplaced: {}, + SessionEventMessageUpdated: {}, + SessionEventMessageInserted: {}, + SessionEventMessagesDeleted: {}, + SessionEventModelContext: {}, + SessionEventRollback: {}, + SessionEventSessionStatusRunning: {}, + SessionEventSessionStatusIdle: {}, + SessionEventSessionError: {}, + SessionEventSpanModelRequestStart: {}, + SessionEventSpanModelRequestEnd: {}, + SessionEventSpanToolCallStart: {}, + SessionEventSpanToolCallEnd: {}, + SessionEventCancel: {}, + SessionEventInterrupt: {}, +} + +func isKnownSessionEventKind(kind SessionEventKind) bool { + if kind == "" { + return false + } + if strings.HasPrefix(string(kind), SessionEventExtensionPrefix) { + return true + } + _, ok := knownSessionEventKinds[kind] + return ok +} + +type LifecycleEvent struct { + State SessionRunState `json:"state,omitempty"` + StopReason *StopReason `json:"stop_reason,omitempty"` +} + +type SessionRollbackEvent struct { + ToEventID string `json:"to_event_id"` + PreviousHeadCommitEventID string `json:"previous_head_commit_event_id,omitempty"` +} + +type SessionRunState string + +const ( + SessionRunStateRunning SessionRunState = "running" + SessionRunStateIdle SessionRunState = "idle" +) + +type StopReason struct { + Type string `json:"type,omitempty"` +} + +type ModelContextEvent struct { + ToolInfos []*schema.ToolInfo `json:"tool_infos,omitempty"` + DeferredToolInfos []*schema.ToolInfo `json:"deferred_tool_infos,omitempty"` +} + +type SessionErrorEvent struct { + // Type identifies the timeline error category. Known values are + // SessionErrorTypeModelRetry, SessionErrorTypeModelFailover, and SessionErrorTypeFatal. + Type string `json:"type,omitempty"` + Message string `json:"message,omitempty"` + RetryStatus *RetryStatus `json:"retry_status,omitempty"` +} + +const ( + SessionErrorTypeModelRetry = "model_retry" + SessionErrorTypeModelFailover = "model_failover" + SessionErrorTypeFatal = "fatal" +) + +type RetryStatus struct { + Type string `json:"type,omitempty"` +} + +type SpanEvent struct { + SpanID string `json:"span_id"` + ParentSpanID string `json:"parent_span_id,omitempty"` + + Kind SpanKind `json:"kind"` + Name string `json:"name,omitempty"` + + StartedAt time.Time `json:"started_at,omitempty"` + EndedAt time.Time `json:"ended_at,omitempty"` + + TTFTMS int64 `json:"ttft_ms,omitempty"` + + Status string `json:"status,omitempty"` + Err string `json:"err,omitempty"` + + // Model and Tool are mutually exclusive: exactly one must be non-nil for + // every Span-carrying SessionEvent. ClassifySessionEvent enforces this + // invariant. + Model *ModelSpanMeta `json:"model,omitempty"` + Tool *ToolSpanMeta `json:"tool,omitempty"` +} + +type SpanKind string + +const ( + SpanKindModel SpanKind = "model" + SpanKindTool SpanKind = "tool" +) + +type ModelSpanMeta struct { + Provider string `json:"provider,omitempty"` + // Model is the model name from options (model.WithModel). Best-effort: empty + // if the user configures model name directly on the ChatModel implementation + // without passing model.WithModel in call-site options. + Model string `json:"model,omitempty"` + Attempt int `json:"attempt,omitempty"` + ModelRequestStartEventID string `json:"model_request_start_event_id,omitempty"` + Usage *ModelUsage `json:"usage,omitempty"` + FinishReason string `json:"finish_reason,omitempty"` + Timeout *ModelTimeoutMeta `json:"timeout,omitempty"` + Accepted bool `json:"accepted"` +} + +// ModelTimeoutMeta records timeout details for a model span that ended with a ModelTimeoutError. +type ModelTimeoutMeta struct { + // Phase identifies which part of the model call exceeded its timeout budget. + Phase string `json:"phase,omitempty"` + // TimeoutMS is the configured timeout budget in milliseconds. + TimeoutMS int64 `json:"timeout_ms,omitempty"` + // ElapsedMS is the observed elapsed duration in milliseconds. + ElapsedMS int64 `json:"elapsed_ms,omitempty"` + // ChunksReceived is the number of stream chunks delivered before the timeout. + ChunksReceived int `json:"chunks_received,omitempty"` +} + +type ModelUsage struct { + InputTokens int `json:"input_tokens,omitempty"` + OutputTokens int `json:"output_tokens,omitempty"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` + Raw *schema.TokenUsage `json:"raw,omitempty"` +} + +// ToolSpanMeta carries the operational metadata of a single tool call span. +// Inputs and outputs are NOT recorded here — they live on the assistant +// message and the tool result message respectively. The span is a stable +// identity envelope that joins those two messages together with timing +// and status. +// +// Tool spans for permission-gated calls may straddle multiple Run/Resume +// invocations: the start span fires on the run where the call begins +// (typically before the user is asked), and the end span fires on the run +// where the call completes (after the user has approved/rejected/responded). +// Both spans share the same SpanID. Consumers correlating a start span to +// its eventual end span should follow SessionEvent.Span.SpanID (or use +// ToolUseID for cross-event correlation across resume boundaries). +type ToolSpanMeta struct { + // ToolUseID is the model-assigned call ID; joins to the assistant + // message's tool-call entry and the tool result message's call ID. + ToolUseID string `json:"tool_use_id"` + + // Name is the tool name. Carried on both start and end so UIs can render + // the span without resolving the assistant message. + Name string `json:"name,omitempty"` + + // ToolCallStartEventID links the end span back to its start (mirrors + // ModelSpanMeta.ModelRequestStartEventID). Set only on the end span. + ToolCallStartEventID string `json:"tool_call_start_event_id,omitempty"` + + // AssistantMessageEventID is the SessionEvent ID of the assistant + // message that emitted this tool call. Lets consumers fetch arguments + // without scanning. Stable across interrupt/resume; the assistant message + // ID established in the original turn is preserved on the eventual end + // span via the in-flight span snapshot. + AssistantMessageEventID string `json:"assistant_message_event_id,omitempty"` + + // ToolResultMessageEventID is the SessionEvent ID of the tool result + // message. Set only on the end span; empty when the call errored before + // producing one. + ToolResultMessageEventID string `json:"tool_result_message_event_id,omitempty"` +} + +// CancelEvent records a user-initiated cancellation in the durable session timeline. +type CancelEvent struct { + Reason string `json:"reason,omitempty"` +} + +// InterruptEvent records a business interrupt in the durable session timeline. +type InterruptEvent struct { + // Contexts is the set of interrupt contexts that caused the agent to pause. + // Each element represents a single root-cause interrupt point. + Contexts []*InterruptContext `json:"contexts,omitempty"` +} + +// InterruptContext describes a single interrupt point within a batch. +type InterruptContext struct { + // InterruptID is the fully-qualified address of the interrupt point + // (e.g. "agent:A;tool:lookup:call_1"). Use this as the key in ResumeParams.Targets. + InterruptID string `json:"interrupt_id,omitempty"` + // Info is the business-defined payload describing the interrupt, provided by + // the component that triggered it (e.g. a middleware or a custom tool). + // ADK treats it as opaque; consumers type-assert it to the concrete type the + // triggering component documents (e.g. *permission.AskInfo) to determine how + // to handle the interrupt. + Info any `json:"info,omitempty"` + // ToolUseID is set when the interrupt source is a specific tool call. It is + // structural metadata derived from the interrupt address, identifying which + // tool call paused; it carries no business semantics. + ToolUseID string `json:"tool_use_id,omitempty"` +} + +// SessionExtensionEvent carries application-owned timeline event payloads. +// The SessionEvent.Kind field is the application event type and must use the +// SessionEventExtensionPrefix namespace. Data is application-owned typed payload +// data. Custom payload types that need durable round-trip behavior must be +// registered with schema.RegisterName before session events are encoded and +// decoded. Consumers can inspect SessionEvent.Kind and type-assert Data to the +// registered concrete payload type. +type SessionExtensionEvent struct { + Data any `json:"data,omitempty"` +} + +// MessageStreamIncompleteEvent records the materialized prefix of a stream that +// failed before EOF. It is durable replay data and does not enter model context. +type MessageStreamIncompleteEvent[M MessageType] struct { + Message M `json:"message"` + Error string `json:"error,omitempty"` +} + +// MessageUpdatedEvent represents a single message replacement within the messages array. +type MessageUpdatedEvent[M MessageType] struct { + // MessageID identifies the target message via its eino-internal message ID + // (stored in Extra["_eino_msg_id"]). UUID v4 assigned by ChatModelAgent for each + // assistant output and tool result, guaranteed unique across turns. + MessageID string `json:"message_id"` + // Message is the new content (with placeholder). + Message M `json:"message"` +} + +// MessageInsertedEvent represents a message inserted by a middleware. +type MessageInsertedEvent[M MessageType] struct { + // Message is the inserted message (carries its own idempotency markers in Extra/metadata). + Message M `json:"message"` + // BeforeMessageID identifies the message BEFORE which this message was inserted, + // using the eino message ID. Empty string means "append at end". + BeforeMessageID string `json:"before_message_id,omitempty"` +} + +// MessagesDeletedEvent represents a batch deletion within the messages array. +type MessagesDeletedEvent struct { + // MessageIDs identifies the messages to delete via their eino-internal message IDs. + MessageIDs []string `json:"message_ids"` +} + +// SessionEventIDGenerator returns the EventID for a draft SessionEvent[M]. +// +// Generators see the fully-populated session-local draft (Kind, Message, Span, +// Extension, payload, timestamp, ...) and may return a business-side identifier +// such as the matching application order/job/result ID. When a generator does not +// recognize a draft event, it should fall through to +// DefaultSessionEventIDGenerator[M] rather than allocating a UUID directly, +// so that the default behavior stays consistent with the framework default. +// +// A returned empty event_id is treated as a generator-side contract violation +// (ErrSessionEventIDGeneratorEmpty); the runner fails closed before the +// event is appended to the store. +type SessionEventIDGenerator[M MessageType] func(ctx context.Context, event *SessionEvent[M]) (string, error) + +// DefaultSessionEventIDGenerator returns a UUID-based EventID for any draft. +// It is exported so that application-side SessionEventIDGenerator[M] +// implementations can fall through to default behavior when they do not +// recognize a draft event, e.g.: +// +// func myGen(ctx context.Context, e *SessionEvent[M]) (string, error) { +// if id, ok := mapDraftToBusinessID(e); ok { +// return id, nil +// } +// return DefaultSessionEventIDGenerator[M](ctx, e) +// } +func DefaultSessionEventIDGenerator[M MessageType](_ context.Context, _ *SessionEvent[M]) (string, error) { + return uuid.NewString(), nil +} + +// SessionConfig tunes managed-session admission and event identity. +type SessionConfig[M MessageType] struct { + // EventIDGenerator decides the EventID of every SessionEvent[M] produced + // by the runner / wrappers. The generator sees the fully-populated draft + // before assignment and may map it to a business-side ID. If nil, + // DefaultSessionEventIDGenerator[M] (UUID v4) is used. + // + // The generator is the sole authority for runner-generated event IDs: + // drafts always have an empty EventID at the assignment boundary, and + // the generator is always invoked. Returning an empty string fails the + // turn closed (ErrSessionEventIDGeneratorEmpty). + EventIDGenerator SessionEventIDGenerator[M] + // SessionAcquireTimeout bounds how long Runner may wait to acquire any session + // handle before failing the current Run/Resume/Rollback attempt. + // + // It applies to the process-local admission path used by the built-in + // session store. + SessionAcquireTimeout time.Duration +} + +type reconstructedSessionState[M MessageType] struct { + Messages []M + ToolInfos []*schema.ToolInfo + DeferredToolInfos []*schema.ToolInfo + sawModelContext bool +} + +type runnerSessionCheckpoint struct { + SessionID string + CheckPointID string + Payload []byte +} + +func init() { + // Register SessionEvent and helper types for HumanReadableSerializer. + schema.RegisterName[*SessionEvent[*schema.Message]]("_eino_adk_session_event") + schema.RegisterName[*SessionEvent[*schema.AgenticMessage]]("_eino_adk_agentic_session_event") + schema.RegisterName[*MessageStreamIncompleteEvent[*schema.Message]]("_eino_adk_message_stream_incomplete_event") + schema.RegisterName[*MessageStreamIncompleteEvent[*schema.AgenticMessage]]("_eino_adk_agentic_message_stream_incomplete_event") + schema.RegisterName[*MessageUpdatedEvent[*schema.Message]]("_eino_adk_message_updated_event") + schema.RegisterName[*MessageUpdatedEvent[*schema.AgenticMessage]]("_eino_adk_agentic_message_updated_event") + schema.RegisterName[*MessageInsertedEvent[*schema.Message]]("_eino_adk_message_inserted_event") + schema.RegisterName[*MessageInsertedEvent[*schema.AgenticMessage]]("_eino_adk_agentic_message_inserted_event") + schema.RegisterName[*MessagesDeletedEvent]("_eino_adk_messages_deleted_event") + schema.RegisterName[*ModelContextEvent]("_eino_adk_model_context_event") + schema.RegisterName[*LifecycleEvent]("_eino_adk_lifecycle_event") + schema.RegisterName[*SessionErrorEvent]("_eino_adk_session_error_event") + schema.RegisterName[*RetryStatus]("_eino_adk_retry_status") + schema.RegisterName[*SpanEvent]("_eino_adk_span_event") + schema.RegisterName[*ModelSpanMeta]("_eino_adk_model_span_meta") + schema.RegisterName[*ModelTimeoutMeta]("_eino_adk_model_timeout_meta") + schema.RegisterName[*ModelUsage]("_eino_adk_model_usage") + schema.RegisterName[*ToolSpanMeta]("_eino_adk_tool_span_meta") + // Note: SessionEventVariant and MessageStreamRef are not registered here + // because they never reach a serializer: checkpoint sanitizers strip + // SessionEventVariant before gob encoding, and the store serializer only + // encodes *SessionEvent[M] (the materialized form, not the variant). + schema.RegisterName[*CancelEvent]("_eino_adk_cancel_event") + schema.RegisterName[*InterruptEvent]("_eino_adk_interrupt_event") + schema.RegisterName[*InterruptContext]("_eino_adk_interrupt_context") + schema.RegisterName[*SessionExtensionEvent]("_eino_adk_session_extension_event") + schema.RegisterName[*SessionRollbackEvent]("_eino_adk_session_rollback_event") +} + +func encodeGob(v any) ([]byte, error) { + var buf bytes.Buffer + if err := gob.NewEncoder(&buf).Encode(v); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func encodeRunnerSessionCheckpoint(c *runnerSessionCheckpoint) ([]byte, error) { + return encodeGob(c) +} + +func decodeRunnerSessionCheckpoint(payload []byte) (*runnerSessionCheckpoint, error) { + var c runnerSessionCheckpoint + if err := gob.NewDecoder(bytes.NewReader(payload)).Decode(&c); err != nil { + return nil, err + } + return &c, nil +} + +func sessionRunnerCheckpointID(sessionID string) string { + return "session/" + sessionID + sessionRunnerCheckpointSuffix +} + +var sessionSerializer schema.Serializer = &schema.HumanReadableSerializer{} + +func encodeSessionEvent[M MessageType](event *SessionEvent[M]) ([]byte, error) { + return encodeSessionEventWithSerializer(event, sessionSerializer) +} + +func decodeSessionEvent[M MessageType](data []byte) (*SessionEvent[M], error) { + return decodeSessionEventWithSerializer[M](data, sessionSerializer) +} + +func encodeSessionEventWithSerializer[M MessageType](event *SessionEvent[M], serializer schema.Serializer) ([]byte, error) { + if err := NormalizeSessionEventKind(event); err != nil { + return nil, err + } + return normalizeSerializer(serializer).Marshal(event) +} + +func decodeSessionEventWithSerializer[M MessageType](data []byte, serializer schema.Serializer) (*SessionEvent[M], error) { + var event SessionEvent[M] + if err := normalizeSerializer(serializer).Unmarshal(data, &event); err != nil { + return nil, err + } + if err := NormalizeSessionEventKind(&event); err != nil { + return nil, err + } + return &event, nil +} + +func snapshotSessionEvent[M MessageType](event *SessionEvent[M]) (*SessionEvent[M], error) { + data, err := encodeSessionEvent(event) + if err != nil { + return nil, err + } + return decodeSessionEvent[M](data) +} + +func normalizeSerializer(serializer schema.Serializer) schema.Serializer { + if serializer == nil { + return sessionSerializer + } + return serializer +} + +// makeInputSessionEvent wraps an input message as a SessionEvent draft. +// +// The returned draft has an empty EventID; the caller must assign one via +// assignSessionEventIDFromContext (or assignSessionEventID) before sending or +// persisting the event. +func makeInputSessionEvent[M MessageType](msg M) *SessionEvent[M] { + return &SessionEvent[M]{Timestamp: newEventTimestamp(), Kind: SessionEventMessage, Message: msg} +} + +// toSessionEvent converts an internal TypedAgentEvent into the persistence format. +// Returns nil if the event has no persistable content. +func toSessionEvent[M MessageType](event *TypedAgentEvent[M]) *SessionEvent[M] { + se, _ := toSessionEventChecked(event) + return se +} + +func toSessionEventChecked[M MessageType](event *TypedAgentEvent[M]) (*SessionEvent[M], error) { + if event == nil { + return nil, nil + } + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil { + se, err := normalizeAgentSessionEvent(event) + if err != nil { + return nil, err + } + if err := ValidateEmittedSessionEventKind(&se); err != nil { + return nil, err + } + if se.Kind == SessionEventModelContext { + return nil, nil + } + return &se, nil + } + if event.Output != nil && event.Output.MessageOutput != nil && + !isNilMessage(event.Output.MessageOutput.Message) { + return nil, errors.New("persistable AgentEvent has no SessionEventVariant.Event") + } + return nil, nil +} + +func normalizeAgentSessionEvent[M MessageType](event *TypedAgentEvent[M]) (SessionEvent[M], error) { + return normalizeAgentSessionEventWithAssigner(event, func(*SessionEvent[M]) (string, error) { + return uuid.NewString(), nil + }) +} + +// normalizeAgentSessionEventWithAssigner unifies the agent event / session +// event identity, allocating a fresh ID via assign when neither side carries +// one. The assigner is invoked with the (still-empty-ID) draft session event +// so callers may route allocation through SessionEventIDGenerator[M]. +func normalizeAgentSessionEventWithAssigner[M MessageType]( + event *TypedAgentEvent[M], + assign func(*SessionEvent[M]) (string, error), +) (SessionEvent[M], error) { + if event == nil || event.SessionEventVariant == nil || event.SessionEventVariant.Event == nil { + return SessionEvent[M]{}, errors.New("missing session event") + } + if assign == nil { + assign = func(*SessionEvent[M]) (string, error) { return uuid.NewString(), nil } + } + se := *event.SessionEventVariant.Event + if se.EventID == "" { + id, err := assign(&se) + if err != nil { + return SessionEvent[M]{}, err + } + if id == "" { + return SessionEvent[M]{}, ErrSessionEventIDGeneratorEmpty + } + se.EventID = id + } + if se.Timestamp.IsZero() { + se.Timestamp = newEventTimestamp() + } + event.SessionEventVariant.Event = &se + return se, nil +} + +func validateAgentSessionEventIdentity[M MessageType](event *TypedAgentEvent[M]) error { + if event == nil || event.SessionEventVariant == nil { + return nil + } + if (event.SessionEventVariant.Event == nil) == (event.SessionEventVariant.MessageStreamRef == nil) { + return errors.New("session event variant must set exactly one payload") + } + if ref := event.SessionEventVariant.MessageStreamRef; ref != nil && ref.Kind != SessionEventMessage { + return fmt.Errorf("message stream ref kind must be %q, got %q", SessionEventMessage, ref.Kind) + } + return nil +} + +// ClassifySessionEvent derives the canonical event kind from the single active +// payload carried by event. +func ClassifySessionEvent[M MessageType](event *SessionEvent[M]) (SessionEventKind, error) { + if event == nil { + return "", errors.New("nil session event") + } + var kinds []SessionEventKind + add := func(kind SessionEventKind) { + kinds = append(kinds, kind) + } + if !isNilMessage(event.Message) { + add(SessionEventMessage) + } + if event.MessageStreamIncomplete != nil { + if isNilMessage(event.MessageStreamIncomplete.Message) { + return "", errors.New("message stream incomplete event must set non-nil Message") + } + add(SessionEventMessageStreamIncomplete) + } + if event.MessagesReplaced != nil { + add(SessionEventMessagesReplaced) + } + if event.MessageUpdated != nil { + add(SessionEventMessageUpdated) + } + if event.MessageInserted != nil { + add(SessionEventMessageInserted) + } + if event.MessagesDeleted != nil { + if err := validateMessageIDs("MessagesDeleted.MessageIDs", event.MessagesDeleted.MessageIDs); err != nil { + return "", err + } + add(SessionEventMessagesDeleted) + } + if event.ModelContext != nil { + if err := validateModelContextEvent(event.ModelContext); err != nil { + return "", err + } + add(SessionEventModelContext) + } + if event.Rollback != nil { + if event.EventID == "" { + return "", errors.New("rollback session event must set non-empty EventID") + } + if event.Rollback.ToEventID == "" { + return "", errors.New("rollback session event must set non-empty ToEventID") + } + add(SessionEventRollback) + } + if event.Lifecycle != nil { + switch event.Lifecycle.State { + case SessionRunStateRunning: + add(SessionEventSessionStatusRunning) + case SessionRunStateIdle: + add(SessionEventSessionStatusIdle) + default: + return "", fmt.Errorf("unknown lifecycle state %q", event.Lifecycle.State) + } + } + if event.Error != nil { + add(SessionEventSessionError) + } + if event.Span != nil { + kind, err := classifySpanSessionEvent(event.Span) + if err != nil { + return "", err + } + add(kind) + } + if event.Cancel != nil { + add(SessionEventCancel) + } + if event.Interrupt != nil { + add(SessionEventInterrupt) + } + if event.Extension != nil { + if event.Kind == "" { + return "", errors.New("session extension event must set kind") + } + if !strings.HasPrefix(string(event.Kind), SessionEventExtensionPrefix) { + return "", fmt.Errorf("session extension event kind %q must start with %q", event.Kind, SessionEventExtensionPrefix) + } + add(event.Kind) + } + if len(kinds) != 1 { + return "", fmt.Errorf("session event must have exactly one active payload, got %d", len(kinds)) + } + return kinds[0], nil +} + +func countActiveSessionEventPayloads[M MessageType](event *SessionEvent[M]) int { + if event == nil { + return 0 + } + count := 0 + if !isNilMessage(event.Message) { + count++ + } + if event.MessageStreamIncomplete != nil { + count++ + } + if event.MessagesReplaced != nil { + count++ + } + if event.MessageUpdated != nil { + count++ + } + if event.MessageInserted != nil { + count++ + } + if event.MessagesDeleted != nil { + count++ + } + if event.ModelContext != nil { + count++ + } + if event.Rollback != nil { + count++ + } + if event.Lifecycle != nil { + count++ + } + if event.Error != nil { + count++ + } + if event.Span != nil { + count++ + } + if event.Cancel != nil { + count++ + } + if event.Interrupt != nil { + count++ + } + if event.Extension != nil { + count++ + } + return count +} + +func classifySpanSessionEvent(span *SpanEvent) (SessionEventKind, error) { + if (span.Model != nil) == (span.Tool != nil) { + return "", errors.New("span event must populate exactly one of Model or Tool") + } + switch span.Kind { + case SpanKindModel: + if span.Model == nil { + return "", errors.New("model span requires Span.Model meta") + } + switch { + case !span.StartedAt.IsZero() && span.EndedAt.IsZero(): + return SessionEventSpanModelRequestStart, nil + case !span.EndedAt.IsZero(): + return SessionEventSpanModelRequestEnd, nil + default: + return "", errors.New("model span must have start or end timestamp") + } + case SpanKindTool: + if span.Tool == nil { + return "", errors.New("tool span requires Span.Tool meta") + } + switch { + case !span.StartedAt.IsZero() && span.EndedAt.IsZero(): + return SessionEventSpanToolCallStart, nil + case !span.EndedAt.IsZero(): + return SessionEventSpanToolCallEnd, nil + default: + return "", errors.New("tool span must have start or end timestamp") + } + default: + return "", fmt.Errorf("unknown span kind %q", span.Kind) + } +} + +// NormalizeSessionEventKind fills an empty Kind from the active payload and +// rejects mismatches between Kind and payload shape. +// +// Unknown kinds (kinds not in the known set and not prefixed with "x.") with +// no recognized payload are tolerated as a forward/backward compatibility +// mechanism: legacy data (e.g. pre-refactor "turn_end") and events written by +// newer code with kinds we don't yet understand are accepted as-is rather than +// failing replay. +func NormalizeSessionEventKind[M MessageType](event *SessionEvent[M]) error { + if event != nil && event.Kind != "" && !isKnownSessionEventKind(event.Kind) && countActiveSessionEventPayloads(event) == 0 { + return nil + } + kind, err := ClassifySessionEvent(event) + if err != nil { + return err + } + if event.Kind != "" && event.Kind != kind { + return fmt.Errorf("session event kind %q does not match payload %q", event.Kind, kind) + } + event.Kind = kind + return nil +} + +// ValidateEmittedSessionEventKind enforces that runtime-emitted session events +// carry an explicit Kind matching their active payload. +func ValidateEmittedSessionEventKind[M MessageType](event *SessionEvent[M]) error { + if event == nil { + return errors.New("nil session event") + } + if event.Kind == "" { + return errors.New("emitted session event must set non-empty Kind") + } + return NormalizeSessionEventKind(event) +} + +func isSessionDurableBoundaryKind(kind SessionEventKind) bool { + switch kind { + case SessionEventMessage, SessionEventSessionStatusIdle, SessionEventInterrupt: + return true + default: + return false + } +} + +func normalizeSessionConfig[M MessageType](cfg *SessionConfig[M]) SessionConfig[M] { + normalized := SessionConfig[M]{ + EventIDGenerator: DefaultSessionEventIDGenerator[M], + SessionAcquireTimeout: defaultSessionAcquireTimeout, + } + if cfg == nil { + return normalized + } + if cfg.EventIDGenerator != nil { + normalized.EventIDGenerator = cfg.EventIDGenerator + } + if cfg.SessionAcquireTimeout > 0 { + normalized.SessionAcquireTimeout = cfg.SessionAcquireTimeout + } + return normalized +} + +// assignSessionEventID assigns the EventID of a draft SessionEvent[M] using +// gen, falling back to DefaultSessionEventIDGenerator[M] when gen is nil. It +// is the single authoritative entry point for SessionEvent[M] ID allocation +// in ADK; runner / wrappers paths must route every draft through this helper +// (or its context wrapper assignSessionEventIDFromContext) before sending or +// persisting the event. +// +// Callers MUST construct the draft with EventID == "" and populate every +// other relevant session-local field (Kind, payload, timestamp) so the +// generator sees a complete draft. A nil event is a no-op. +// +// On generator-side contract violations, the helper returns: +// - ErrSessionEventIDGeneratorEmpty when gen returns an empty id; +// - the generator's wrapped error otherwise. +// +// The runner is expected to fail closed on these errors and not append the +// event to the store. +func assignSessionEventID[M MessageType]( + ctx context.Context, + event *SessionEvent[M], + gen SessionEventIDGenerator[M], +) error { + if event == nil { + return nil + } + if gen == nil { + gen = DefaultSessionEventIDGenerator[M] + } + id, err := gen(ctx, event) + if err != nil { + return fmt.Errorf("adk: session event id generator: %w", err) + } + if id == "" { + return ErrSessionEventIDGeneratorEmpty + } + event.EventID = id + return nil +} + +type sessionEventIDGeneratorKey[M MessageType] struct{} + +// contextWithSessionEventIDGenerator stores the typed SessionEventIDGenerator[M] +// in ctx so that deeply-nested wrappers (model / tool / middleware) can route +// SessionEvent[M] draft ID allocation through the runner's configured +// generator without explicit parameter threading. +// +// This is an internal plumbing escape hatch; the only payload allowed in ctx +// under this key is the generator function itself. Storing business IDs, +// per-event state, or anything else under this key is forbidden — see the +// "scoped exception" note in the design plan. +func contextWithSessionEventIDGenerator[M MessageType](ctx context.Context, gen SessionEventIDGenerator[M]) context.Context { + if gen == nil { + return ctx + } + return context.WithValue(ctx, sessionEventIDGeneratorKey[M]{}, gen) +} + +func sessionEventIDGeneratorFromContext[M MessageType](ctx context.Context) SessionEventIDGenerator[M] { + if v := ctx.Value(sessionEventIDGeneratorKey[M]{}); v != nil { + if gen, ok := v.(SessionEventIDGenerator[M]); ok { + return gen + } + } + return nil +} + +// assignSessionEventIDFromContext assigns the EventID of a draft +// SessionEvent[M] using the SessionEventIDGenerator[M] stored in ctx (falling +// back to DefaultSessionEventIDGenerator[M] when none is set). Wrappers and +// runner closures call this helper after populating the rest of the draft. +func assignSessionEventIDFromContext[M MessageType](ctx context.Context, event *SessionEvent[M]) error { + if event == nil { + return nil + } + return assignSessionEventID(ctx, event, sessionEventIDGeneratorFromContext[M](ctx)) +} + +type sessionEventPersister[M MessageType] struct { + ctx context.Context + handle sessionHandle[M] + sessionID string + pending []*SessionEvent[M] + + mu sync.Mutex + err error +} + +func newSessionEventPersister[M MessageType]( + ctx context.Context, + handle sessionHandle[M], + sessionID string, +) *sessionEventPersister[M] { + return &sessionEventPersister[M]{ + ctx: ctx, + handle: handle, + sessionID: sessionID, + } +} + +func (p *sessionEventPersister[M]) enqueueAsync(event *SessionEvent[M]) error { + if event == nil || event.EventID == "" { + return p.getErr() + } + snapshot, err := snapshotSessionEvent(event) + if err != nil { + p.setErr(err) + return err + } + if err := p.getErr(); err != nil { + return err + } + p.mu.Lock() + p.pending = append(p.pending, snapshot) + p.mu.Unlock() + return nil +} + +func (p *sessionEventPersister[M]) commitBoundary(event *SessionEvent[M]) error { + if event == nil || event.EventID == "" { + return p.getErr() + } + snapshot, err := snapshotSessionEvent(event) + if err != nil { + p.setErr(err) + return err + } + if err := p.flushPending(); err != nil { + return err + } + return p.appendEvents([]*SessionEvent[M]{snapshot}) +} + +func (p *sessionEventPersister[M]) flushPending() error { + if err := p.getErr(); err != nil { + return err + } + p.mu.Lock() + events := make([]*SessionEvent[M], len(p.pending)) + copy(events, p.pending) + p.mu.Unlock() + if len(events) == 0 { + return nil + } + if err := p.appendEvents(events); err != nil { + return err + } + p.mu.Lock() + p.pending = nil + p.mu.Unlock() + return nil +} + +func (p *sessionEventPersister[M]) closeAndWait() error { + return p.flushPending() +} + +func (p *sessionEventPersister[M]) appendEvents(events []*SessionEvent[M]) error { + err := p.handle.appendEvents(p.ctx, events) + if err != nil { + p.setErr(err) + return err + } + return nil +} + +func (p *sessionEventPersister[M]) setErr(err error) { + if err == nil { + return + } + p.mu.Lock() + if p.err == nil { + p.err = err + } + p.mu.Unlock() +} + +func (p *sessionEventPersister[M]) getErr() error { + p.mu.Lock() + defer p.mu.Unlock() + return p.err +} + +func stripSessionEventFields[M MessageType](event *TypedAgentEvent[M]) *TypedAgentEvent[M] { + if event == nil { + return nil + } + if event.SessionEventVariant == nil { + return event + } + stripped := *event + stripped.SessionEventVariant = nil + if stripped.Output == nil && stripped.Action == nil && stripped.Err == nil { + return nil + } + return &stripped +} + +// applySessionEvent applies a single SessionEvent to the message array, mutating in place. +// Non-message events are ignored. +func applySessionEvent[M MessageType](messages *[]M, event *SessionEvent[M]) error { + if !isContextSessionEvent(event) { + return nil + } + return applyContextSessionEventInPlace(event, messages) +} + +func isContextSessionEvent[M MessageType](event *SessionEvent[M]) bool { + if event == nil { + return false + } + return !isNilMessage(event.Message) || event.MessagesReplaced != nil || + event.MessageUpdated != nil || event.MessageInserted != nil || event.MessagesDeleted != nil +} + +func applyContextSessionEvent[M MessageType](messages []M, event *SessionEvent[M]) ([]M, error) { + out := append([]M{}, messages...) + err := applyContextSessionEventInPlace(event, &out) + return out, err +} + +func applyContextSessionEventInPlace[M MessageType](event *SessionEvent[M], out *[]M) error { + switch { + case event.MessagesReplaced != nil: + *out = append([]M{}, *event.MessagesReplaced...) + + case event.MessageUpdated != nil: + upd := event.MessageUpdated + if replacementID := GetMessageID(upd.Message); replacementID != "" && replacementID != upd.MessageID { + return fmt.Errorf("apply event: MessageUpdated target %q but replacement has ID %q — identity mismatch", upd.MessageID, replacementID) + } + if err := replaceMessageByID(out, upd.MessageID, upd.Message); err != nil { + return err + } + + case event.MessageInserted != nil: + ins := event.MessageInserted + if ins.BeforeMessageID == "" { + *out = append(*out, ins.Message) + } else { + inserted := false + for j, msg := range *out { + if GetMessageID(msg) == ins.BeforeMessageID { + var zero M + *out = append(*out, zero) + copy((*out)[j+1:], (*out)[j:]) + (*out)[j] = ins.Message + inserted = true + break + } + } + if !inserted { + return fmt.Errorf("apply event: anchor message %q not found for insertion", ins.BeforeMessageID) + } + } + + case event.MessagesDeleted != nil: + if err := deleteMessagesByID(out, event.MessagesDeleted.MessageIDs); err != nil { + return err + } + + default: + if !isNilMessage(event.Message) { + *out = append(*out, event.Message) + } + } + return nil +} + +// replaceMessageByID finds the message with the given ID and replaces it. +func replaceMessageByID[M MessageType](messages *[]M, msgID string, newMsg M) error { + for i, msg := range *messages { + if GetMessageID(msg) == msgID { + (*messages)[i] = newMsg + return nil + } + } + return fmt.Errorf("reconstruct: target message %q not found for update", msgID) +} + +func deleteMessagesByID[M MessageType](messages *[]M, ids []string) error { + if err := validateMessageIDs("MessagesDeleted.MessageIDs", ids); err != nil { + return err + } + targets := make(map[string]struct{}, len(ids)) + for _, id := range ids { + targets[id] = struct{}{} + } + found := make(map[string]struct{}, len(ids)) + for _, msg := range *messages { + id := GetMessageID(msg) + if _, ok := targets[id]; ok { + found[id] = struct{}{} + } + } + for _, id := range ids { + if _, ok := found[id]; !ok { + return fmt.Errorf("reconstruct: target message %q not found for deletion", id) + } + } + + retained := (*messages)[:0] + for _, msg := range *messages { + if _, ok := targets[GetMessageID(msg)]; ok { + continue + } + retained = append(retained, msg) + } + *messages = retained + return nil +} + +func validateMessageIDs(field string, ids []string) error { + if len(ids) == 0 { + return fmt.Errorf("%s must not be empty", field) + } + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + if id == "" { + return fmt.Errorf("%s must not contain empty message ID", field) + } + if _, ok := seen[id]; ok { + return fmt.Errorf("%s contains duplicate message ID %q", field, id) + } + seen[id] = struct{}{} + } + return nil +} + +func validateModelContextEvent(event *ModelContextEvent) error { + if event == nil { + return nil + } + if err := validateToolInfoNames("ModelContext.ToolInfos", event.ToolInfos); err != nil { + return err + } + return validateToolInfoNames("ModelContext.DeferredToolInfos", event.DeferredToolInfos) +} + +func validateToolInfoNames(field string, infos []*schema.ToolInfo) error { + seen := make(map[string]struct{}, len(infos)) + for _, info := range infos { + if info == nil { + continue + } + if info.Name == "" { + return fmt.Errorf("%s must not contain empty tool name", field) + } + if _, ok := seen[info.Name]; ok { + return fmt.Errorf("%s contains duplicate tool name %q", field, info.Name) + } + seen[info.Name] = struct{}{} + } + return nil +} + +type sessionReconstructResult[M MessageType] struct { + state *reconstructedSessionState[M] +} + +// sessionReplayEventKinds is the set of event kinds required to project the active +// session log and reconstruct model-facing state. +var sessionReplayEventKinds = []SessionEventKind{ + SessionEventMessage, + SessionEventMessagesReplaced, + SessionEventMessageUpdated, + SessionEventMessageInserted, + SessionEventMessagesDeleted, + SessionEventModelContext, + SessionEventSessionStatusIdle, + SessionEventInterrupt, + SessionEventCancel, + SessionEventRollback, +} + +type RollbackSessionOptions[M MessageType] struct { + CheckPointStore CheckPointStore + ExpectedHeadEventID string + EventIDGenerator SessionEventIDGenerator[M] +} + +type RollbackSessionOption[M MessageType] func(*RollbackSessionOptions[M]) + +// WithRollbackSessionCheckPointStore deletes session-derived checkpoints after a successful rollback. +func WithRollbackSessionCheckPointStore[M MessageType](store CheckPointStore) RollbackSessionOption[M] { + return func(opts *RollbackSessionOptions[M]) { + opts.CheckPointStore = store + } +} + +// WithRollbackSessionExpectedHeadEventID requires the current active committed +// idle event to match eventID before rollback. +func WithRollbackSessionExpectedHeadEventID[M MessageType](eventID string) RollbackSessionOption[M] { + return func(opts *RollbackSessionOptions[M]) { + opts.ExpectedHeadEventID = eventID + } +} + +// WithRollbackEventIDGenerator overrides the EventID generator for the rollback +// event. The generator sees the fully-populated rollback draft (kind, +// SessionRollbackEvent payload, timestamp) before assignment. If nil or not set, +// DefaultSessionEventIDGenerator[M] (UUID v4) is used. +func WithRollbackEventIDGenerator[M MessageType](gen SessionEventIDGenerator[M]) RollbackSessionOption[M] { + return func(opts *RollbackSessionOptions[M]) { + opts.EventIDGenerator = gen + } +} + +// RollbackSession appends a rollback marker that makes the committed idle event +// with targetEventID the latest active committed boundary. +func RollbackSession[M MessageType]( + ctx context.Context, + store SessionEventStore[M], + sessionID string, + targetEventID string, + opts ...RollbackSessionOption[M], +) error { + if store == nil { + return errors.New("adk: rollback session store is nil") + } + if sessionID == "" { + return errors.New("adk: rollback sessionID is empty") + } + if targetEventID == "" { + return ErrRollbackTargetNotFound + } + var cfg RollbackSessionOptions[M] + for _, opt := range opts { + if opt != nil { + opt(&cfg) + } + } + openResult, err := openRunnerSession[M](ctx, store, sessionID, normalizeSessionConfig[M](nil)) + if err != nil { + return err + } + if openResult == nil || openResult.handle == nil { + return ErrSessionBusy + } + defer openResult.handle.close(ctx) + + activeEvents, err := loadActiveSessionEventsReverse[M](ctx, openResult.handle, sessionID, defaultLoadPageSize) + if err != nil { + return err + } + target, head, err := resolveRollbackTarget[M](activeEvents, targetEventID) + if err != nil { + if errors.Is(err, ErrRollbackTargetNotFound) { + evidence, evidenceErr := findPhysicalRollbackTargetEvidence[M](ctx, openResult.handle, sessionID, targetEventID, defaultLoadPageSize) + if evidenceErr != nil { + return evidenceErr + } + switch evidence { + case rollbackTargetEvidenceCommitted: + return ErrRollbackTargetInactive + case rollbackTargetEvidenceUncommitted: + return ErrInvalidRollbackTarget + } + } + return err + } + if cfg.ExpectedHeadEventID != "" && (head == nil || head.EventID != cfg.ExpectedHeadEventID) { + return ErrSessionHeadChanged + } + + rb := &SessionEvent[M]{ + Timestamp: newEventTimestamp(), + Kind: SessionEventRollback, + Rollback: &SessionRollbackEvent{ + ToEventID: target.EventID, + PreviousHeadCommitEventID: head.EventID, + }, + } + if err := assignSessionEventID(ctx, rb, cfg.EventIDGenerator); err != nil { + return err + } + if err := openResult.handle.appendEvents(ctx, []*SessionEvent[M]{rb}); err != nil { + return err + } + if cfg.CheckPointStore != nil { + if deleter, ok := cfg.CheckPointStore.(CheckPointDeleter); ok { + if err := deleter.Delete(ctx, sessionRunnerCheckpointID(sessionID)); err != nil { + return fmt.Errorf("failed to delete session checkpoint after rollback: %w", err) + } + } + } + return nil +} + +// reconstructSessionState rebuilds session state by replaying the active log. +// Committed idle lifecycle events are loaded for rollback projection but are not +// applied as reconstructed state. +func reconstructSessionState[M MessageType]( + ctx context.Context, + handle sessionHandle[M], + sessionID string, + pageSize int, +) (*sessionReconstructResult[M], error) { + allEvents, err := loadActiveSessionEventsReverse[M](ctx, handle, sessionID, pageSize) + if err != nil { + return nil, err + } + if len(allEvents) == 0 { + return nil, nil + } + + state, err := replayDurableContextEvents(allEvents) + if err != nil { + return nil, err + } + return &sessionReconstructResult[M]{state: state}, nil +} + +func loadActiveSessionEventsReverse[M MessageType]( + ctx context.Context, + handle sessionHandle[M], + sessionID string, + pageSize int, +) ([]*SessionEvent[M], error) { + if pageSize <= 0 { + pageSize = defaultLoadPageSize + } + var physicalReverse []*SessionEvent[M] + var after string + for { + result, err := handle.loadEvents(ctx, &LoadSessionEventsRequest{ + After: after, + Limit: pageSize, + Reverse: true, + Kinds: sessionReplayEventKinds, + }) + if err != nil { + return nil, err + } + if result == nil || len(result.Events) == 0 { + break + } + physicalReverse = append(physicalReverse, result.Events...) + if result.Next == "" { + break + } + after = result.Next + } + return projectActiveEventsFromReverse(physicalReverse) +} + +func projectActiveEventsFromReverse[M MessageType]( + physicalReverse []*SessionEvent[M], +) ([]*SessionEvent[M], error) { + active := make([]*SessionEvent[M], 0, len(physicalReverse)) + activeLen := 0 + posByEventID := make(map[string]int, len(physicalReverse)) + for i := len(physicalReverse) - 1; i >= 0; i-- { + event := physicalReverse[i] + if event.Kind == SessionEventRollback { + rb, err := decodeRollbackSessionEvent(event) + if err != nil { + return nil, err + } + pos, ok := posByEventID[rb.ToEventID] + if !ok || pos >= activeLen || active[pos].EventID != rb.ToEventID { + return nil, ErrRollbackTargetInactive + } + target := active[pos] + if !isCommittedIdleEvent(target) { + return nil, ErrInvalidRollbackTarget + } + activeLen = pos + 1 + continue + } + if activeLen < len(active) { + active[activeLen] = event + active = active[:activeLen+1] + } else { + active = append(active, event) + } + posByEventID[event.EventID] = activeLen + activeLen++ + } + if activeLen < len(active) { + active = active[:activeLen] + } + return active, nil +} + +type rollbackTargetEvidence int + +const ( + rollbackTargetEvidenceNone rollbackTargetEvidence = iota + rollbackTargetEvidenceUncommitted + rollbackTargetEvidenceCommitted +) + +func findPhysicalRollbackTargetEvidence[M MessageType]( + ctx context.Context, + handle sessionHandle[M], + sessionID string, + targetEventID string, + pageSize int, +) (rollbackTargetEvidence, error) { + if pageSize <= 0 { + pageSize = defaultLoadPageSize + } + var after string + var evidence rollbackTargetEvidence + for { + result, err := handle.loadEvents(ctx, &LoadSessionEventsRequest{ + After: after, + Limit: pageSize, + Reverse: false, + Kinds: sessionReplayEventKinds, + }) + if err != nil { + return rollbackTargetEvidenceNone, err + } + if result == nil || len(result.Events) == 0 { + break + } + for _, event := range result.Events { + if event.Kind == SessionEventRollback { + continue + } + if event.EventID != targetEventID { + continue + } + if isCommittedIdleEvent(event) { + return rollbackTargetEvidenceCommitted, nil + } + evidence = rollbackTargetEvidenceUncommitted + } + if result.Next == "" { + break + } + after = result.Next + } + return evidence, nil +} + +func decodeRollbackSessionEvent[M MessageType](event *SessionEvent[M]) (*SessionRollbackEvent, error) { + if event == nil || event.EventID == "" || event.Kind != SessionEventRollback || event.Rollback == nil { + return nil, ErrInvalidRollbackTarget + } + if event.Rollback.ToEventID == "" { + return nil, ErrInvalidRollbackTarget + } + return event.Rollback, nil +} + +func resolveRollbackTarget[M MessageType]( + activeEvents []*SessionEvent[M], + targetEventID string, +) (target *SessionEvent[M], head *SessionEvent[M], err error) { + var sawTargetEvent bool + for _, event := range activeEvents { + if event != nil && event.EventID == targetEventID { + sawTargetEvent = true + } + if !isCommittedIdleEvent(event) { + continue + } + head = event + if event.EventID == targetEventID { + target = event + } + } + if target != nil { + return target, head, nil + } + if sawTargetEvent { + return nil, nil, ErrInvalidRollbackTarget + } + return nil, nil, ErrRollbackTargetNotFound +} + +func replayDurableContextEvents[M MessageType](events []*SessionEvent[M]) (*reconstructedSessionState[M], error) { + if len(events) == 0 { + return nil, nil + } + var messages []M + startIdx := 0 + boundaryIdx := -1 + for i := 0; i < len(events); i++ { + if events[i].MessagesReplaced != nil { + boundaryIdx = i + } + } + + if boundaryIdx >= 0 { + messages = append([]M{}, *events[boundaryIdx].MessagesReplaced...) + startIdx = boundaryIdx + 1 + } + + // Model-context reconstruction is intentionally scoped to events at or after + // the latest MessagesReplaced boundary, the same window used for messages. + // A model_context emitted before that boundary is not recovered: doing so + // would require scanning the full session log, which defeats the point of + // shortcutting reconstruction at MessagesReplaced. The only consequence is + // that the next turn's first model call re-emits a model_context snapshot + // (sawModelContext starts false) even when the tool set is unchanged — an + // extra audit event, not a correctness issue, since reconstructed ToolInfos + // feed only the change-detection baseline and never the model itself. + state := &reconstructedSessionState[M]{Messages: messages} + for i := startIdx; i < len(events); i++ { + if err := applySessionEvent(&messages, events[i]); err != nil { + return nil, fmt.Errorf("reconstruct: %w", err) + } + if events[i] != nil && events[i].ModelContext != nil { + state.ToolInfos = cloneToolInfos(events[i].ModelContext.ToolInfos) + state.DeferredToolInfos = cloneToolInfos(events[i].ModelContext.DeferredToolInfos) + state.sawModelContext = true + } + } + state.Messages = messages + return state, nil +} + +func isCommittedIdleEvent[M MessageType](event *SessionEvent[M]) bool { + return event != nil && + event.Kind == SessionEventSessionStatusIdle && + event.Lifecycle != nil && + event.Lifecycle.State == SessionRunStateIdle && + event.Lifecycle.StopReason != nil && + event.Lifecycle.StopReason.Type == "end_turn" +} diff --git a/adk/session/conformance.go b/adk/session/conformance.go new file mode 100644 index 000000000..a2bba0c6d --- /dev/null +++ b/adk/session/conformance.go @@ -0,0 +1,440 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package session provides session event stores and a reusable conformance test +// suite for validating SessionEventStore implementations. +package session + +import ( + "context" + "errors" + "fmt" + "reflect" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +type conformanceExtensionPayload struct { + OK bool `json:"ok"` +} + +func init() { + schema.RegisterName[*conformanceExtensionPayload]("_eino_adk_session_conformance_extension_payload") +} + +// RunConformanceTests validates the SessionEventStore contract shared by +// provider-facing session persistence implementations. +// +// The contract assumes single-writer-per-session: tests do NOT exercise +// concurrent AppendEvents calls for the same sessionID. +func RunConformanceTests[M adk.MessageType]( + t *testing.T, + factory func(testing.TB) adk.SessionEventStore[M], + makeMessage func(content string) M, +) { + t.Helper() + + t.Run("AppendEvents and forward LoadEvents", func(t *testing.T) { testAppendAndForwardLoad(t, factory, makeMessage) }) + t.Run("LoadEvents reverse pagination", func(t *testing.T) { testReversePagination(t, factory, makeMessage) }) + t.Run("After forward pagination", func(t *testing.T) { testForwardPagination(t, factory, makeMessage) }) + t.Run("sessionID isolates events", func(t *testing.T) { testSessionIsolation(t, factory, makeMessage) }) + t.Run("Empty session returns no events", func(t *testing.T) { testEmptySession(t, factory) }) + t.Run("AppendEvents rejects non-replay duplicate EventID", func(t *testing.T) { testRejectDuplicateEventID(t, factory, makeMessage) }) + t.Run("AppendEvents rejects duplicate EventID within same batch", func(t *testing.T) { testRejectDuplicateEventIDWithinBatch(t, factory, makeMessage) }) + t.Run("AppendEvents rejects empty EventID with ErrInvalidEventID", func(t *testing.T) { testRejectEmptyEventID(t, factory, makeMessage) }) + t.Run("After resumes by EventID forward", func(t *testing.T) { testAfterForward(t, factory, makeMessage) }) + t.Run("After resumes by EventID reverse", func(t *testing.T) { testAfterReverse(t, factory, makeMessage) }) + t.Run("Unknown After returns ErrEventIDOutOfRange", func(t *testing.T) { testUnknownAfter(t, factory, makeMessage) }) + t.Run("Empty page when After=last forward and After=first reverse", func(t *testing.T) { testEmptyPageBoundary(t, factory, makeMessage) }) + t.Run("Extension kind filters correctly", func(t *testing.T) { testExtensionKindFilter(t, factory) }) + t.Run("event body round-trips", func(t *testing.T) { testEventBodyRoundTrip(t, factory, makeMessage) }) +} + +// RunSerializerConformanceTests validates that a concrete SessionEventStore +// implementation honors its implementation-local serializer configuration. +func RunSerializerConformanceTests[M adk.MessageType]( + t *testing.T, + factory func(testing.TB, schema.Serializer) adk.SessionEventStore[M], + makeMessage func(content string) M, +) { + t.Helper() + t.Run("custom serializer is honored", func(t *testing.T) { + serializer := &countingEventSerializer{inner: &schema.HumanReadableSerializer{}} + store := factory(t, serializer) + if store == nil { + t.Fatalf("factory returned nil SessionEventStore") + } + + ctx := context.Background() + event := messageEvent("custom-serializer-1", makeMessage("custom serializer")) + appendEvents(t, ctx, store, "s", event) + if serializer.marshalCount == 0 { + t.Fatalf("custom serializer Marshal was not called") + } + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + if serializer.unmarshalCount == 0 { + t.Fatalf("custom serializer Unmarshal was not called") + } + requireEventsEqual(t, []*adk.SessionEvent[M]{event}, res.Events) + }) +} + +func testAppendAndForwardLoad[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + first := messageEvent("e1", makeMessage("first")) + second := committedIdleEvent[M]("e2") + third := messageEvent("e3", makeMessage("third")) + appendEvents(t, ctx, store, "s", first, second) + appendEvents(t, ctx, store, "s", third) + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + if res == nil { + t.Fatalf("LoadEvents returned nil result") + } + requireEventsEqual(t, []*adk.SessionEvent[M]{first, second, third}, res.Events) +} + +func testExtensionKindFilter[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M]) { + store := newStore(t, factory) + ctx := context.Background() + + first := extensionEvent[M]("custom-1", "x.conformance.custom") + second := committedIdleEvent[M]("turn-1") + third := extensionEvent[M]("custom-2", "x.conformance.custom") + appendEvents(t, ctx, store, "s", first, second, third) + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{ + Kinds: []adk.SessionEventKind{adk.SessionEventKind("x.conformance.custom")}, + }) + requireNoError(t, err) + if res == nil { + t.Fatalf("LoadEvents returned nil result") + } + requireEventsEqual(t, []*adk.SessionEvent[M]{first, third}, res.Events) +} + +func testReversePagination[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + events := make([]*adk.SessionEvent[M], 5) + for i := 0; i < 5; i++ { + events[i] = messageEvent(fmt.Sprintf("r%d", i), makeMessage(fmt.Sprintf("%c", 'a'+i))) + appendEvents(t, ctx, store, "s", events[i]) + } + + var collected []string + var after string + for { + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{ + Reverse: true, + Limit: 2, + After: after, + }) + requireNoError(t, err) + if res == nil || len(res.Events) == 0 { + break + } + for _, ep := range res.Events { + collected = append(collected, ep.EventID) + } + if res.Next == "" { + break + } + after = res.Next + } + + expected := []string{"r4", "r3", "r2", "r1", "r0"} + if len(collected) != len(expected) { + t.Fatalf("reverse collected length=%d want=%d (got=%v)", len(collected), len(expected), collected) + } + for i := range expected { + if collected[i] != expected[i] { + t.Fatalf("reverse[%d]=%q want=%q (got=%v)", i, collected[i], expected[i], collected) + } + } +} + +func testForwardPagination[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + for i := 0; i < 80; i++ { + event := messageEvent(fmt.Sprintf("f%d", i), makeMessage(fmt.Sprintf("%d", i))) + appendEvents(t, ctx, store, "s", event) + } + + var collected []*adk.SessionEvent[M] + req := &adk.LoadSessionEventsRequest{Limit: 10} + for { + res, err := store.LoadEvents(ctx, "s", req) + requireNoError(t, err) + if res == nil || len(res.Events) == 0 { + break + } + collected = append(collected, res.Events...) + if res.Next == "" { + break + } + req = &adk.LoadSessionEventsRequest{Limit: 10, After: res.Next} + } + if len(collected) != 80 { + t.Fatalf("expected 80 events, got %d", len(collected)) + } + for i, ep := range collected { + expectedID := fmt.Sprintf("f%d", i) + if ep.EventID != expectedID { + t.Fatalf("event[%d].EventID=%q, want=%q", i, ep.EventID, expectedID) + } + } +} + +func testSessionIsolation[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + alpha := messageEvent("alpha-1", makeMessage("alpha")) + beta := committedIdleEvent[M]("beta-1") + appendEvents(t, ctx, store, "alpha", alpha) + appendEvents(t, ctx, store, "beta", beta) + + alphaRes, err := store.LoadEvents(ctx, "alpha", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + requireEventsEqual(t, []*adk.SessionEvent[M]{alpha}, alphaRes.Events) + + betaRes, err := store.LoadEvents(ctx, "beta", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + requireEventsEqual(t, []*adk.SessionEvent[M]{beta}, betaRes.Events) +} + +func testEmptySession[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M]) { + store := newStore(t, factory) + ctx := context.Background() + + res, err := store.LoadEvents(ctx, "nonexistent", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + if res != nil && len(res.Events) != 0 { + t.Fatalf("expected empty result for nonexistent session, got %d events", len(res.Events)) + } +} + +func testRejectDuplicateEventID[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + first := messageEvent("dup-1", makeMessage("first")) + dup := messageEvent("dup-1", makeMessage("second")) + appendEvents(t, ctx, store, "s", first) + err := store.AppendEvents(ctx, "s", []*adk.SessionEvent[M]{dup}) + if !errors.Is(err, adk.ErrDuplicateEventID) { + t.Fatalf("expected ErrDuplicateEventID, got %v", err) + } + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + requireEventsEqual(t, []*adk.SessionEvent[M]{first}, res.Events) +} + +func testRejectDuplicateEventIDWithinBatch[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + first := messageEvent("dup-batch-1", makeMessage("first")) + dup := messageEvent("dup-batch-1", makeMessage("second")) + err := store.AppendEvents(ctx, "s", []*adk.SessionEvent[M]{first, dup}) + if !errors.Is(err, adk.ErrDuplicateEventID) { + t.Fatalf("expected ErrDuplicateEventID, got %v", err) + } + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + requireEventsEqual(t, nil, res.Events) +} + +func testRejectEmptyEventID[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + err := store.AppendEvents(ctx, "s", []*adk.SessionEvent[M]{{Kind: adk.SessionEventMessage, Message: makeMessage("empty")}}) + if !errors.Is(err, adk.ErrInvalidEventID) { + t.Fatalf("expected ErrInvalidEventID, got %v", err) + } +} + +func testAfterForward[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + events := make([]*adk.SessionEvent[M], 5) + for i := 0; i < 5; i++ { + events[i] = messageEvent(fmt.Sprintf("fwd-%d", i), makeMessage(fmt.Sprintf("%d", i))) + appendEvents(t, ctx, store, "s", events[i]) + } + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{After: "fwd-2"}) + requireNoError(t, err) + requireEventsEqual(t, []*adk.SessionEvent[M]{events[3], events[4]}, res.Events) +} + +func testAfterReverse[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + events := make([]*adk.SessionEvent[M], 5) + for i := 0; i < 5; i++ { + events[i] = messageEvent(fmt.Sprintf("rev-%d", i), makeMessage(fmt.Sprintf("%d", i))) + appendEvents(t, ctx, store, "s", events[i]) + } + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{Reverse: true, After: "rev-2"}) + requireNoError(t, err) + requireEventsEqual(t, []*adk.SessionEvent[M]{events[1], events[0]}, res.Events) +} + +func testUnknownAfter[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + appendEvents(t, ctx, store, "s", messageEvent("only-1", makeMessage("only"))) + + _, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{After: "ghost"}) + if !errors.Is(err, adk.ErrEventIDOutOfRange) { + t.Fatalf("forward unknown After expected ErrEventIDOutOfRange, got %v", err) + } + _, err = store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{After: "ghost", Reverse: true}) + if !errors.Is(err, adk.ErrEventIDOutOfRange) { + t.Fatalf("reverse unknown After expected ErrEventIDOutOfRange, got %v", err) + } +} + +func testEmptyPageBoundary[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + ids := []string{"e0", "e1", "e2"} + for _, id := range ids { + appendEvents(t, ctx, store, "s", messageEvent(id, makeMessage(id))) + } + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{After: "e2"}) + requireNoError(t, err) + if res == nil || len(res.Events) != 0 || res.Next != "" { + t.Fatalf("forward empty page expected, got events=%d next=%q", len(res.Events), res.Next) + } + + res, err = store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{Reverse: true, After: "e0"}) + requireNoError(t, err) + if res == nil || len(res.Events) != 0 || res.Next != "" { + t.Fatalf("reverse empty page expected, got events=%d next=%q", len(res.Events), res.Next) + } +} + +func testEventBodyRoundTrip[M adk.MessageType](t *testing.T, factory func(testing.TB) adk.SessionEventStore[M], makeMessage func(string) M) { + store := newStore(t, factory) + ctx := context.Background() + + event := messageEvent("body-test-1", makeMessage("body")) + appendEvents(t, ctx, store, "s", event) + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + requireNoError(t, err) + if res == nil || len(res.Events) != 1 { + t.Fatalf("expected 1 event, got %d", len(res.Events)) + } + requireEventsEqual(t, []*adk.SessionEvent[M]{event}, res.Events) +} + +func newStore[M adk.MessageType](t testing.TB, factory func(testing.TB) adk.SessionEventStore[M]) adk.SessionEventStore[M] { + t.Helper() + store := factory(t) + if store == nil { + t.Fatalf("factory returned nil SessionEventStore") + } + return store +} + +func appendEvents[M adk.MessageType](t testing.TB, ctx context.Context, store adk.SessionEventStore[M], sessionID string, events ...*adk.SessionEvent[M]) { + t.Helper() + err := store.AppendEvents(ctx, sessionID, events) + requireNoError(t, err) +} + +func requireNoError(t testing.TB, err error) { + t.Helper() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func requireEventsEqual[M adk.MessageType](t testing.TB, want, got []*adk.SessionEvent[M]) { + t.Helper() + if len(want) != len(got) { + t.Fatalf("events length mismatch: got=%d want=%d", len(got), len(want)) + } + for i := range want { + if !reflect.DeepEqual(got[i], want[i]) { + t.Fatalf("event[%d] mismatch:\n got: %#v\nwant: %#v", i, got[i], want[i]) + } + } +} + +type countingEventSerializer struct { + inner schema.Serializer + marshalCount int + unmarshalCount int +} + +func (s *countingEventSerializer) Marshal(v any) ([]byte, error) { + s.marshalCount++ + return s.inner.Marshal(v) +} + +func (s *countingEventSerializer) Unmarshal(data []byte, v any) error { + s.unmarshalCount++ + return s.inner.Unmarshal(data, v) +} + +func messageEvent[M adk.MessageType](id string, msg M) *adk.SessionEvent[M] { + return &adk.SessionEvent[M]{EventID: id, Kind: adk.SessionEventMessage, Message: msg} +} + +func committedIdleEvent[M adk.MessageType](id string) *adk.SessionEvent[M] { + return &adk.SessionEvent[M]{ + EventID: id, + Kind: adk.SessionEventSessionStatusIdle, + Lifecycle: &adk.LifecycleEvent{ + State: adk.SessionRunStateIdle, + StopReason: &adk.StopReason{Type: "end_turn"}, + }, + } +} + +func extensionEvent[M adk.MessageType](id, kind string) *adk.SessionEvent[M] { + return &adk.SessionEvent[M]{ + EventID: id, + Kind: adk.SessionEventKind(kind), + Extension: &adk.SessionExtensionEvent{ + Data: &conformanceExtensionPayload{OK: true}, + }, + } +} diff --git a/adk/session/file_store.go b/adk/session/file_store.go new file mode 100644 index 000000000..2ac7d2fd3 --- /dev/null +++ b/adk/session/file_store.go @@ -0,0 +1,452 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package session + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// FileStoreConfig configures FileStore. +type FileStoreConfig struct { + // EventSerializer encodes typed session events before storage. Defaults to + // schema.HumanReadableSerializer. Output must not contain raw CR/LF bytes. + EventSerializer schema.Serializer +} + +// FileStore is a process-local, file-backed implementation of adk.SessionEventStore. +// Each session is stored as one event log file under the configured directory: +// +// /.evlog +// +// Each line is formatted as: \t\t\n +// where Data is the raw serialized bytes written directly to the line. +// +// IMPORTANT: FileStore requires that serialized event data does NOT contain raw +// newline (\n) or carriage-return (\r) characters, because these would +// corrupt the line-oriented file format. The default HumanReadableSerializer +// (compact JSON) satisfies this constraint. Serializers that may emit \n or \r +// in their output (e.g. GobSerializer, raw protobuf) are NOT compatible with +// FileStore — use InMemoryStore or a custom store implementation instead. +// AppendEvents will return an error if Data contains \n or \r. +// +// FileStore does not implement CheckPointStore; runner checkpoints should use a +// dedicated checkpoint store. +// +// FileStore synchronizes access within the current process. It does not provide +// cross-process write safety. +type FileStore[M adk.MessageType] struct { + dir string + serializer schema.Serializer + mu sync.Mutex + indexes map[string]*fileSessionIndex +} + +type fileEvent struct { + eventID string + kind adk.SessionEventKind + data []byte +} + +type fileSessionIndex struct { + size int64 + modTime time.Time + offsets []int64 + eventIDToLine map[string]int +} + +// NewFileStore creates a file-backed SessionEventStore rooted at dir. +func NewFileStore[M adk.MessageType](dir string, cfg *FileStoreConfig) (*FileStore[M], error) { + if dir == "" { + return nil, errorsNewEmptyFileStoreDir() + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, err + } + return &FileStore[M]{ + dir: dir, + serializer: normalizeFileSerializer(cfg), + indexes: make(map[string]*fileSessionIndex), + }, nil +} + +func errorsNewEmptyFileStoreDir() error { + return fmt.Errorf("adk/session: file store dir is empty") +} + +func errorsNewEmptySessionID() error { + return fmt.Errorf("adk/session: sessionID is empty") +} + +// AppendEvents appends events to the session's event log. +// +// Each SessionEvent.EventID MUST be non-empty. Duplicate event IDs are rejected. +func (s *FileStore[M]) AppendEvents(_ context.Context, sessionID string, events []*adk.SessionEvent[M]) error { + s.mu.Lock() + defer s.mu.Unlock() + + path, err := s.sessionPath(sessionID) + if err != nil { + return err + } + + // Validate incoming events and dedup within batch. + seen := make(map[string]struct{}, len(events)) + pending := make([]fileEvent, 0, len(events)) + for _, e := range events { + if e == nil || e.EventID == "" { + return adk.ErrInvalidEventID + } + if _, dup := seen[e.EventID]; dup { + return adk.ErrDuplicateEventID + } + seen[e.EventID] = struct{}{} + if normalizeErr := adk.NormalizeSessionEventKind(e); normalizeErr != nil { + return normalizeErr + } + data, marshalErr := s.serializer.Marshal(e) + if marshalErr != nil { + return marshalErr + } + if bytes.ContainsAny(data, "\r\n") { + return fmt.Errorf("adk/session: FileStore requires serialized event data without raw CR/LF; use a line-safe serializer") + } + pending = append(pending, fileEvent{eventID: e.EventID, kind: e.Kind, data: data}) + } + if len(pending) == 0 { + return nil + } + + idx, err := s.ensureIndexLocked(path) + if err != nil { + return err + } + + var out *os.File + for _, event := range pending { + if _, dup := idx.eventIDToLine[event.eventID]; dup { + return adk.ErrDuplicateEventID + } + if out == nil { + out, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644) + if err != nil { + return err + } + defer out.Close() + } + line := fmt.Sprintf("%s\t%s\t%s\n", event.eventID, event.kind, event.data) + n, err := out.WriteString(line) + if err != nil { + return err + } + idx.eventIDToLine[event.eventID] = len(idx.offsets) + idx.offsets = append(idx.offsets, idx.size) + idx.size += int64(n) + } + if out != nil { + info, err := out.Stat() + if err != nil { + return err + } + idx.size = info.Size() + idx.modTime = info.ModTime() + } + return nil +} + +// LoadEvents loads events with pagination and direction support. +func (s *FileStore[M]) LoadEvents(_ context.Context, sessionID string, opts *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[M], error) { + s.mu.Lock() + defer s.mu.Unlock() + if opts == nil { + opts = &adk.LoadSessionEventsRequest{} + } + path, err := s.sessionPath(sessionID) + if err != nil { + return nil, err + } + idx, err := s.ensureIndexLocked(path) + if err != nil { + return nil, err + } + if opts.Reverse { + return s.loadFileEventsReverseLocked(path, idx, opts) + } + return s.loadFileEventsForwardLocked(path, idx, opts) +} + +func (s *FileStore[M]) sessionPath(sessionID string) (string, error) { + if sessionID == "" { + return "", errorsNewEmptySessionID() + } + return filepath.Join(s.dir, url.PathEscape(sessionID)+".evlog"), nil +} + +func (s *FileStore[M]) ensureIndexLocked(path string) (*fileSessionIndex, error) { + info, err := os.Stat(path) + if err != nil { + if os.IsNotExist(err) { + idx := &fileSessionIndex{eventIDToLine: make(map[string]int)} + s.indexes[path] = idx + return idx, nil + } + return nil, err + } + if idx := s.indexes[path]; idx != nil && idx.size == info.Size() && idx.modTime.Equal(info.ModTime()) { + return idx, nil + } + idx, err := s.rebuildIndexLocked(path, info) + if err != nil { + delete(s.indexes, path) + return nil, err + } + s.indexes[path] = idx + return idx, nil +} + +func (s *FileStore[M]) rebuildIndexLocked(path string, info os.FileInfo) (*fileSessionIndex, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + reader := bufio.NewReader(f) + idx := &fileSessionIndex{ + size: info.Size(), + modTime: info.ModTime(), + eventIDToLine: make(map[string]int), + } + lineNo := 0 + var offset int64 + for { + lineOffset := offset + line, readErr := reader.ReadBytes('\n') + if len(line) > 0 { + offset += int64(len(line)) + lineNo++ + event, err := parseFileEventLine(line, lineNo) + if err != nil { + return nil, err + } + if _, dup := idx.eventIDToLine[event.eventID]; dup { + return nil, fmt.Errorf("%w: duplicate event_id %q at line %d", adk.ErrInvalidEventID, event.eventID, lineNo) + } + idx.eventIDToLine[event.eventID] = len(idx.offsets) + idx.offsets = append(idx.offsets, lineOffset) + } + if readErr == nil { + continue + } + if readErr == io.EOF { + break + } + return nil, readErr + } + return idx, nil +} + +func parseFileEventLine(line []byte, lineNo int) (fileEvent, error) { + if len(line) == 0 || line[len(line)-1] != '\n' { + return fileEvent{}, fmt.Errorf("%w: corrupted trailing record at line %d", adk.ErrInvalidEventID, lineNo) + } + line = line[:len(line)-1] + lineStr := string(line) + + firstTab := strings.IndexByte(lineStr, '\t') + if firstTab < 0 { + return fileEvent{}, fmt.Errorf("%w: missing tab separator at line %d", adk.ErrInvalidEventID, lineNo) + } + eventID := lineStr[:firstTab] + if eventID == "" { + return fileEvent{}, fmt.Errorf("%w: empty event_id at line %d", adk.ErrInvalidEventID, lineNo) + } + + rest := lineStr[firstTab+1:] + secondTab := strings.IndexByte(rest, '\t') + if secondTab < 0 { + return fileEvent{}, fmt.Errorf("%w: missing kind tab separator at line %d", adk.ErrInvalidEventID, lineNo) + } + return fileEvent{ + eventID: eventID, + kind: adk.SessionEventKind(rest[:secondTab]), + data: []byte(rest[secondTab+1:]), + }, nil +} + +func (s *FileStore[M]) loadFileEventsForwardLocked(path string, idx *fileSessionIndex, opts *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[M], error) { + start := 0 + if opts.After != "" { + pos, ok := idx.eventIDToLine[opts.After] + if !ok { + return nil, adk.ErrEventIDOutOfRange + } + start = pos + 1 + } + if start > len(idx.offsets) { + start = len(idx.offsets) + } + + f, err := os.Open(path) + if err != nil { + if os.IsNotExist(err) { + return &adk.LoadSessionEventsResult[M]{}, nil + } + return nil, err + } + defer f.Close() + kindSet := buildKindSet(opts.Kinds) + + var out []*adk.SessionEvent[M] + hasMore := false + for i := start; i < len(idx.offsets); i++ { + event, err := readFileEventAt(f, idx.offsets[i], i+1) + if err != nil { + return nil, err + } + if kindSet != nil { + if _, match := kindSet[event.kind]; !match { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + decoded, err := s.decodeFileEvent(event) + if err != nil { + return nil, err + } + out = append(out, decoded) + } + + var next string + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &adk.LoadSessionEventsResult[M]{Events: out, Next: next}, nil +} + +func (s *FileStore[M]) loadFileEventsReverseLocked(path string, idx *fileSessionIndex, opts *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[M], error) { + end := len(idx.offsets) + if opts.After != "" { + pos, ok := idx.eventIDToLine[opts.After] + if !ok { + return nil, adk.ErrEventIDOutOfRange + } + end = pos + } + if end <= 0 { + return &adk.LoadSessionEventsResult[M]{}, nil + } + + f, err := os.Open(path) + if err != nil { + if os.IsNotExist(err) { + return &adk.LoadSessionEventsResult[M]{}, nil + } + return nil, err + } + defer f.Close() + kindSet := buildKindSet(opts.Kinds) + + var out []*adk.SessionEvent[M] + hasMore := false + for i := end - 1; i >= 0; i-- { + event, err := readFileEventAt(f, idx.offsets[i], i+1) + if err != nil { + return nil, err + } + if kindSet != nil { + if _, match := kindSet[event.kind]; !match { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + decoded, err := s.decodeFileEvent(event) + if err != nil { + return nil, err + } + out = append(out, decoded) + } + + var next string + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &adk.LoadSessionEventsResult[M]{Events: out, Next: next}, nil +} + +func fileCurrentTailLocked(idx *fileSessionIndex) string { + if idx == nil || len(idx.offsets) == 0 { + return "" + } + for id, line := range idx.eventIDToLine { + if line == len(idx.offsets)-1 { + return id + } + } + return "" +} + +func readFileEventAt(f *os.File, offset int64, lineNo int) (fileEvent, error) { + if _, err := f.Seek(offset, io.SeekStart); err != nil { + return fileEvent{}, err + } + reader := bufio.NewReader(f) + line, err := reader.ReadBytes('\n') + if err != nil { + return fileEvent{}, err + } + return parseFileEventLine(line, lineNo) +} + +func (s *FileStore[M]) decodeFileEvent(src fileEvent) (*adk.SessionEvent[M], error) { + var event adk.SessionEvent[M] + if err := s.serializer.Unmarshal(src.data, &event); err != nil { + return nil, err + } + if err := adk.NormalizeSessionEventKind(&event); err != nil { + return nil, err + } + if event.EventID != src.eventID || event.Kind != src.kind { + return nil, fmt.Errorf("adk/session: file event metadata mismatch for event_id %q", src.eventID) + } + return &event, nil +} + +func normalizeFileSerializer(cfg *FileStoreConfig) schema.Serializer { + if cfg != nil && cfg.EventSerializer != nil { + return cfg.EventSerializer + } + return &schema.HumanReadableSerializer{} +} diff --git a/adk/session/file_store_test.go b/adk/session/file_store_test.go new file mode 100644 index 000000000..17f1683fd --- /dev/null +++ b/adk/session/file_store_test.go @@ -0,0 +1,291 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package session_test + +import ( + "context" + "errors" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/session" + "github.com/cloudwego/eino/schema" +) + +func TestFileStoreConformance(t *testing.T) { + session.RunConformanceTests[*schema.Message](t, func(t testing.TB) adk.SessionEventStore[*schema.Message] { + store, err := session.NewFileStore[*schema.Message](t.TempDir(), nil) + require.NoError(t, err) + return store + }, func(content string) *schema.Message { + return schema.UserMessage(content) + }) + session.RunSerializerConformanceTests[*schema.Message](t, func(t testing.TB, serializer schema.Serializer) adk.SessionEventStore[*schema.Message] { + store, err := session.NewFileStore[*schema.Message](t.TempDir(), &session.FileStoreConfig{EventSerializer: serializer}) + require.NoError(t, err) + return store + }, func(content string) *schema.Message { + return schema.UserMessage(content) + }) +} + +func TestFileStorePersistsAcrossInstances(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + + first := testMessageEvent("persist-1", "first") + second := testCommittedIdleEvent("persist-2", "turn-1") + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{first, second}) + require.NoError(t, err) + + reopened, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + res, err := reopened.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + require.Len(t, res.Events, 2) + assert.Equal(t, "persist-1", res.Events[0].EventID) + assert.Equal(t, "persist-2", res.Events[1].EventID) +} + +func TestFileStoreWritesHumanReadableEvlogLines(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + + first := testMessageEvent("line-1", "first") + second := testCommittedIdleEvent("line-2", "turn-1") + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{first, second}) + require.NoError(t, err) + + data, err := os.ReadFile(filepath.Join(dir, url.PathEscape("s")+".evlog")) + require.NoError(t, err) + lines := strings.Split(strings.TrimSuffix(string(data), "\n"), "\n") + require.Len(t, lines, 2) + + parts0 := strings.SplitN(lines[0], "\t", 3) + require.Len(t, parts0, 3) + assert.Equal(t, "line-1", parts0[0]) + assert.Equal(t, "message", parts0[1]) + assert.Contains(t, parts0[2], "first") + + parts1 := strings.SplitN(lines[1], "\t", 3) + require.Len(t, parts1, 3) + assert.Equal(t, "line-2", parts1[0]) + assert.Equal(t, "session.status_idle", parts1[1]) +} + +func TestFileStoreRollbackPreservesPhysicalAuditLog(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + sessionID := "rollback-audit" + + err = store.AppendEvents(ctx, sessionID, []*adk.SessionEvent[*schema.Message]{ + testMessageEvent("msg-1", "Q1"), + testCommittedIdleEvent("end-1", "turn-1"), + testMessageEvent("msg-2", "Q2"), + testCommittedIdleEvent("end-2", "turn-2"), + }) + require.NoError(t, err) + + require.NoError(t, adk.RollbackSession[*schema.Message](ctx, store, sessionID, "end-1")) + + res, err := store.LoadEvents(ctx, sessionID, &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + require.Len(t, res.Events, 5) + assert.Equal(t, "msg-2", res.Events[2].EventID) + assert.Equal(t, "end-2", res.Events[3].EventID) + assert.Equal(t, adk.SessionEventRollback, res.Events[4].Kind) + + data, err := os.ReadFile(filepath.Join(dir, url.PathEscape(sessionID)+".evlog")) + require.NoError(t, err) + lines := strings.Split(strings.TrimSuffix(string(data), "\n"), "\n") + require.Len(t, lines, 5) + assert.Contains(t, lines[4], "\trollback\t") +} + +func TestFileStoreRejectsInvalidDir(t *testing.T) { + store, err := session.NewFileStore[*schema.Message]("", nil) + require.Error(t, err) + assert.Nil(t, store) +} + +func TestFileStoreRejectsSerializerRawLineDelimiters(t *testing.T) { + ctx := context.Background() + store, err := session.NewFileStore[*schema.Message](t.TempDir(), &session.FileStoreConfig{ + EventSerializer: newlineSerializer{}, + }) + require.NoError(t, err) + + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{testMessageEvent("bad", "bad")}) + require.Error(t, err) + assert.Contains(t, err.Error(), "without raw CR/LF") +} + +func TestFileStoreAppendFailsOnCorruptedExistingLog(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + + path := filepath.Join(dir, url.PathEscape("s")+".evlog") + require.NoError(t, os.WriteFile(path, []byte("corrupted-no-tab\n"), 0o644)) + + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{testMessageEvent("new", "new")}) + require.Error(t, err) + assert.True(t, errors.Is(err, adk.ErrInvalidEventID)) +} + +func TestFileStoreEscapedSessionIDPath(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + + sessionID := "a/b %snow" + err = store.AppendEvents(ctx, sessionID, []*adk.SessionEvent[*schema.Message]{testMessageEvent("escaped", "ok")}) + require.NoError(t, err) + + res, err := store.LoadEvents(ctx, sessionID, &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + require.Len(t, res.Events, 1) + assert.Equal(t, "escaped", res.Events[0].EventID) + + entries, err := os.ReadDir(dir) + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, url.PathEscape(sessionID)+".evlog", entries[0].Name()) +} + +func TestFileStoreValidationReplayAndReversePagination(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + store, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + + _, err = session.NewFileStore[*schema.Message]("", nil) + require.Error(t, err) + + service, err := session.NewFileStore[*schema.Message](filepath.Join(dir, "svc"), nil) + require.NoError(t, err) + assert.NotNil(t, service) + + require.Error(t, store.AppendEvents(ctx, "", nil)) + + empty, err := store.LoadEvents(ctx, "empty", &adk.LoadSessionEventsRequest{Reverse: true}) + require.NoError(t, err) + assert.Empty(t, empty.Events) + + events := []*adk.SessionEvent[*schema.Message]{ + testMessageEvent("e1", "one"), + testSpanEvent("e2"), + testCommittedIdleEvent("e3", "turn-1"), + } + err = store.AppendEvents(ctx, "s", events) + require.NoError(t, err) + + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{testMessageEvent("e4", "four")}) + require.NoError(t, err) + + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{testMessageEvent("e1", "duplicate existing")}) + require.ErrorIs(t, err, adk.ErrDuplicateEventID) + + err = store.AppendEvents(ctx, "s2", []*adk.SessionEvent[*schema.Message]{ + testMessageEvent("dup", "one"), + testMessageEvent("dup", "two"), + }) + require.ErrorIs(t, err, adk.ErrDuplicateEventID) + + _, err = store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{After: "missing"}) + require.ErrorIs(t, err, adk.ErrEventIDOutOfRange) + _, err = store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{Reverse: true, After: "missing"}) + require.ErrorIs(t, err, adk.ErrEventIDOutOfRange) + + forward, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{ + After: "e1", + Kinds: []adk.SessionEventKind{adk.SessionEventSessionStatusIdle, adk.SessionEventMessage}, + Limit: 1, + }) + require.NoError(t, err) + require.Len(t, forward.Events, 1) + assert.Equal(t, "e3", forward.Events[0].EventID) + assert.Equal(t, "e3", forward.Next) + + reverse, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{ + Reverse: true, + After: "e4", + Limit: 1, + }) + require.NoError(t, err) + require.Len(t, reverse.Events, 1) + assert.Equal(t, "e3", reverse.Events[0].EventID) + assert.Equal(t, "e3", reverse.Next) +} + +func TestFileStoreRejectsCorruptedRecordsOnIndexRebuild(t *testing.T) { + ctx := context.Background() + cases := map[string]string{ + "missing newline": "e1\tmessage\t{}", + "empty event id": "\tmessage\t{}\n", + "missing kind tab": "e1\tmessage-only\n", + "duplicate event id": "e1\tmessage\t{}\ne1\tmessage\t{}\n", + "metadata mismatches": "e1\tturn_end\t{\"event_id\":\"e1\",\"kind\":\"message\",\"message\":{\"role\":\"user\",\"content\":\"x\"}}\n", + "invalid event body": "e1\tmessage\tnot-json\n", + "invalid event shape": "e1\tmessage\t{\"event_id\":\"e1\",\"kind\":\"message\"}\n", + "empty session id": "", + } + + for name, content := range cases { + t.Run(name, func(t *testing.T) { + dir := t.TempDir() + store, err := session.NewFileStore[*schema.Message](dir, nil) + require.NoError(t, err) + + if name == "empty session id" { + _, err = store.LoadEvents(ctx, "", &adk.LoadSessionEventsRequest{}) + require.Error(t, err) + return + } + + path := filepath.Join(dir, url.PathEscape("s")+".evlog") + require.NoError(t, os.WriteFile(path, []byte(content), 0o644)) + _, err = store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + require.Error(t, err) + }) + } +} + +type newlineSerializer struct{} + +func (newlineSerializer) Marshal(any) ([]byte, error) { + return []byte("bad\nline"), nil +} + +func (newlineSerializer) Unmarshal([]byte, any) error { + return nil +} diff --git a/adk/session/in_memory_store.go b/adk/session/in_memory_store.go new file mode 100644 index 000000000..187b6301b --- /dev/null +++ b/adk/session/in_memory_store.go @@ -0,0 +1,279 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package session + +import ( + "context" + "fmt" + "sync" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// InMemoryStoreConfig configures InMemoryStore. +type InMemoryStoreConfig struct { + // EventSerializer encodes typed session events before storage. Defaults to + // schema.HumanReadableSerializer. + EventSerializer schema.Serializer +} + +// InMemoryStore is a thread-safe, in-memory implementation of adk.SessionEventStore +// and CheckPointStore (with Delete support). Suitable for testing and +// single-process deployments where durability is not required. +type InMemoryStore[M adk.MessageType] struct { + mu sync.Mutex + events map[string][][]byte + eventIDs map[string][]string + eventKinds map[string][]adk.SessionEventKind + eventIDIdx map[string]map[string]int + serializer schema.Serializer + checkpoints map[string][]byte +} + +type pendingEvent struct { + eventID string + kind adk.SessionEventKind + data []byte +} + +// NewInMemoryStore creates a new InMemoryStore. +func NewInMemoryStore[M adk.MessageType](cfg *InMemoryStoreConfig) *InMemoryStore[M] { + return &InMemoryStore[M]{ + events: make(map[string][][]byte), + eventIDs: make(map[string][]string), + eventKinds: make(map[string][]adk.SessionEventKind), + eventIDIdx: make(map[string]map[string]int), + serializer: normalizeSerializer(cfg), + checkpoints: make(map[string][]byte), + } +} + +// AppendEvents appends events to the session's event log. +func (s *InMemoryStore[M]) AppendEvents(_ context.Context, sessionID string, events []*adk.SessionEvent[M]) error { + s.mu.Lock() + defer s.mu.Unlock() + idx, ok := s.eventIDIdx[sessionID] + if !ok { + idx = make(map[string]int) + s.eventIDIdx[sessionID] = idx + } + seen := make(map[string]struct{}, len(events)) + pending := make([]pendingEvent, 0, len(events)) + for _, e := range events { + if e == nil || e.EventID == "" { + return adk.ErrInvalidEventID + } + if _, dup := seen[e.EventID]; dup { + return adk.ErrDuplicateEventID + } + seen[e.EventID] = struct{}{} + if _, dup := idx[e.EventID]; dup { + return adk.ErrDuplicateEventID + } + if err := adk.NormalizeSessionEventKind(e); err != nil { + return err + } + data, err := s.serializer.Marshal(e) + if err != nil { + return err + } + pending = append(pending, pendingEvent{ + eventID: e.EventID, + kind: e.Kind, + data: append([]byte{}, data...), + }) + } + for _, event := range pending { + s.events[sessionID] = append(s.events[sessionID], event.data) + s.eventIDs[sessionID] = append(s.eventIDs[sessionID], event.eventID) + s.eventKinds[sessionID] = append(s.eventKinds[sessionID], event.kind) + idx[event.eventID] = len(s.events[sessionID]) - 1 + } + return nil +} + +// LoadEvents loads events with pagination and direction support. +func (s *InMemoryStore[M]) LoadEvents(_ context.Context, sessionID string, opts *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[M], error) { + s.mu.Lock() + defer s.mu.Unlock() + + if opts == nil { + opts = &adk.LoadSessionEventsRequest{} + } + if opts.Reverse { + return s.loadReverse(sessionID, opts) + } + return s.loadForward(sessionID, opts) +} + +func (s *InMemoryStore[M]) loadForward(sessionID string, opts *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[M], error) { + all := s.events[sessionID] + idx := s.eventIDIdx[sessionID] + kinds := s.eventKinds[sessionID] + + start := 0 + if opts.After != "" { + pos, ok := idx[opts.After] + if !ok { + return nil, adk.ErrEventIDOutOfRange + } + start = pos + 1 + } + if start > len(all) { + start = len(all) + } + + kindSet := buildKindSet(opts.Kinds) + + var out []*adk.SessionEvent[M] + hasMore := false + for i := start; i < len(all); i++ { + if kindSet != nil { + if _, match := kindSet[kinds[i]]; !match { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + event, err := s.decodeEvent(all[i], s.eventIDs[sessionID][i], kinds[i]) + if err != nil { + return nil, err + } + out = append(out, event) + } + + var next string + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &adk.LoadSessionEventsResult[M]{Events: out, Next: next}, nil +} + +func (s *InMemoryStore[M]) loadReverse(sessionID string, opts *adk.LoadSessionEventsRequest) (*adk.LoadSessionEventsResult[M], error) { + all := s.events[sessionID] + idx := s.eventIDIdx[sessionID] + kinds := s.eventKinds[sessionID] + + end := len(all) + if opts.After != "" { + pos, ok := idx[opts.After] + if !ok { + return nil, adk.ErrEventIDOutOfRange + } + end = pos // strictly older: [0, pos) + } + if end <= 0 { + return &adk.LoadSessionEventsResult[M]{}, nil + } + + kindSet := buildKindSet(opts.Kinds) + + var out []*adk.SessionEvent[M] + hasMore := false + for i := end - 1; i >= 0; i-- { + if kindSet != nil { + if _, match := kindSet[kinds[i]]; !match { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + event, err := s.decodeEvent(all[i], s.eventIDs[sessionID][i], kinds[i]) + if err != nil { + return nil, err + } + out = append(out, event) + } + + var next string + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &adk.LoadSessionEventsResult[M]{Events: out, Next: next}, nil +} + +func (s *InMemoryStore[M]) currentTailLocked(sessionID string) string { + ids := s.eventIDs[sessionID] + if len(ids) == 0 { + return "" + } + return ids[len(ids)-1] +} + +func (s *InMemoryStore[M]) decodeEvent(data []byte, eventID string, kind adk.SessionEventKind) (*adk.SessionEvent[M], error) { + var event adk.SessionEvent[M] + if err := s.serializer.Unmarshal(data, &event); err != nil { + return nil, err + } + if err := adk.NormalizeSessionEventKind(&event); err != nil { + return nil, err + } + if event.EventID != eventID || event.Kind != kind { + return nil, fmt.Errorf("adk/session: in-memory event index mismatch for event_id %q", eventID) + } + return &event, nil +} + +func normalizeSerializer(cfg *InMemoryStoreConfig) schema.Serializer { + if cfg != nil && cfg.EventSerializer != nil { + return cfg.EventSerializer + } + return &schema.HumanReadableSerializer{} +} + +func buildKindSet(kinds []adk.SessionEventKind) map[adk.SessionEventKind]struct{} { + if len(kinds) == 0 { + return nil + } + set := make(map[adk.SessionEventKind]struct{}, len(kinds)) + for _, k := range kinds { + set[k] = struct{}{} + } + return set +} + +// Set stores a checkpoint value. +func (s *InMemoryStore[M]) Set(_ context.Context, checkPointID string, checkPoint []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + s.checkpoints[checkPointID] = append([]byte{}, checkPoint...) + return nil +} + +// Get retrieves a checkpoint value. Returns an independent copy. +func (s *InMemoryStore[M]) Get(_ context.Context, checkPointID string) ([]byte, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.checkpoints[checkPointID] + if !ok { + return nil, false, nil + } + return append([]byte{}, v...), true, nil +} + +// Delete removes a checkpoint. +func (s *InMemoryStore[M]) Delete(_ context.Context, checkPointID string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.checkpoints, checkPointID) + return nil +} diff --git a/adk/session/in_memory_store_test.go b/adk/session/in_memory_store_test.go new file mode 100644 index 000000000..cf5c0de15 --- /dev/null +++ b/adk/session/in_memory_store_test.go @@ -0,0 +1,192 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package session_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/session" + "github.com/cloudwego/eino/schema" +) + +func TestInMemoryStoreConformance(t *testing.T) { + session.RunConformanceTests[*schema.Message](t, func(testing.TB) adk.SessionEventStore[*schema.Message] { + return session.NewInMemoryStore[*schema.Message](nil) + }, func(content string) *schema.Message { + return schema.UserMessage(content) + }) + session.RunSerializerConformanceTests[*schema.Message](t, func(_ testing.TB, serializer schema.Serializer) adk.SessionEventStore[*schema.Message] { + return session.NewInMemoryStore[*schema.Message](&session.InMemoryStoreConfig{EventSerializer: serializer}) + }, func(content string) *schema.Message { + return schema.UserMessage(content) + }) +} + +func TestInMemoryStoreCheckpointSetGetDelete(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + + _, exists, err := store.Get(ctx, "missing") + require.NoError(t, err) + assert.False(t, exists) + + require.NoError(t, store.Set(ctx, "k", []byte("payload"))) + + got, exists, err := store.Get(ctx, "k") + require.NoError(t, err) + require.True(t, exists) + assert.Equal(t, []byte("payload"), got) + + got[0] = 'X' + again, _, err := store.Get(ctx, "k") + require.NoError(t, err) + assert.Equal(t, []byte("payload"), again, "Get must return an independent copy") + + require.NoError(t, store.Delete(ctx, "k")) + _, exists, err = store.Get(ctx, "k") + require.NoError(t, err) + assert.False(t, exists) +} + +func TestInMemoryStoreKindFilterAndPagination(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + events := []*adk.SessionEvent[*schema.Message]{ + testMessageEvent("e1", "one"), + testSpanEvent("e2"), + testCommittedIdleEvent("e3", "turn-1"), + testMessageEvent("e4", "four"), + } + err := store.AppendEvents(ctx, "s", events) + require.NoError(t, err) + + res, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{ + After: "e2", + Kinds: []adk.SessionEventKind{adk.SessionEventMessage, adk.SessionEventSessionStatusIdle}, + Limit: 1, + }) + require.NoError(t, err) + require.Len(t, res.Events, 1) + assert.Equal(t, "e3", res.Events[0].EventID) + assert.Equal(t, "e3", res.Next) +} + +func TestInMemoryStoreLoadReturnsIndependentEvents(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + err := store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{ + testMessageEvent("e1", "one"), + }) + require.NoError(t, err) + + first, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + first.Events[0].EventID = "mutated" + + second, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{}) + require.NoError(t, err) + assert.Equal(t, "e1", second.Events[0].EventID) +} + +func TestInMemoryStoreValidationReplayAndReversePagination(t *testing.T) { + ctx := context.Background() + store := session.NewInMemoryStore[*schema.Message](nil) + + require.NoError(t, store.AppendEvents(ctx, "", nil)) + + events := []*adk.SessionEvent[*schema.Message]{ + testMessageEvent("e1", "one"), + testSpanEvent("e2"), + testCommittedIdleEvent("e3", "turn-1"), + } + err := store.AppendEvents(ctx, "s", events) + require.NoError(t, err) + + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{testMessageEvent("e4", "four")}) + require.NoError(t, err) + + err = store.AppendEvents(ctx, "s2", []*adk.SessionEvent[*schema.Message]{nil}) + require.ErrorIs(t, err, adk.ErrInvalidEventID) + + err = store.AppendEvents(ctx, "s2", []*adk.SessionEvent[*schema.Message]{ + testMessageEvent("dup", "one"), + testMessageEvent("dup", "two"), + }) + require.ErrorIs(t, err, adk.ErrDuplicateEventID) + + err = store.AppendEvents(ctx, "s", []*adk.SessionEvent[*schema.Message]{testMessageEvent("e1", "duplicate existing")}) + require.ErrorIs(t, err, adk.ErrDuplicateEventID) + + err = store.AppendEvents(ctx, "s2", []*adk.SessionEvent[*schema.Message]{{EventID: "invalid-kind"}}) + require.Error(t, err) + + reverseEmpty, err := store.LoadEvents(ctx, "empty", &adk.LoadSessionEventsRequest{Reverse: true}) + require.NoError(t, err) + assert.Empty(t, reverseEmpty.Events) + + _, err = store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{After: "missing"}) + require.ErrorIs(t, err, adk.ErrEventIDOutOfRange) + _, err = store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{Reverse: true, After: "missing"}) + require.ErrorIs(t, err, adk.ErrEventIDOutOfRange) + + reverse, err := store.LoadEvents(ctx, "s", &adk.LoadSessionEventsRequest{ + Reverse: true, + After: "e4", + Limit: 1, + }) + require.NoError(t, err) + require.Len(t, reverse.Events, 1) + assert.Equal(t, "e3", reverse.Events[0].EventID) + assert.Equal(t, "e3", reverse.Next) +} + +func testMessageEvent(id, content string) *adk.SessionEvent[*schema.Message] { + return &adk.SessionEvent[*schema.Message]{ + EventID: id, + Kind: adk.SessionEventMessage, + Message: schema.UserMessage(content), + } +} + +func testCommittedIdleEvent(id, turnID string) *adk.SessionEvent[*schema.Message] { + return &adk.SessionEvent[*schema.Message]{ + EventID: id, + Kind: adk.SessionEventSessionStatusIdle, + Lifecycle: &adk.LifecycleEvent{ + State: adk.SessionRunStateIdle, + StopReason: &adk.StopReason{Type: "end_turn"}, + }, + } +} + +func testSpanEvent(id string) *adk.SessionEvent[*schema.Message] { + return &adk.SessionEvent[*schema.Message]{ + EventID: id, + Kind: adk.SessionEventSpanModelRequestStart, + Span: &adk.SpanEvent{ + Kind: adk.SpanKindModel, + StartedAt: time.Now(), + Model: &adk.ModelSpanMeta{}, + }, + } +} diff --git a/adk/session_admission.go b/adk/session_admission.go new file mode 100644 index 000000000..1c96a29b5 --- /dev/null +++ b/adk/session_admission.go @@ -0,0 +1,116 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "fmt" + "reflect" + "sync" +) + +var localSessionAdmission = struct { + mu sync.Mutex + locked map[string]bool +}{locked: make(map[string]bool)} + +func openLocalSession[M MessageType](_ context.Context, store SessionEventStore[M], req *openSessionRequest) (*openSessionResult[M], error) { + if store == nil || req == nil || req.sessionID == "" { + return nil, ErrSessionBusy + } + key := localSessionAdmissionKey(store, req.sessionID) + if key == "" { + return nil, ErrSessionBusy + } + localSessionAdmission.mu.Lock() + defer localSessionAdmission.mu.Unlock() + if localSessionAdmission.locked[key] { + return nil, ErrSessionBusy + } + localSessionAdmission.locked[key] = true + return &openSessionResult[M]{ + handle: &localSessionHandle[M]{ + key: key, + store: store, + sessionID: req.sessionID, + }, + }, nil +} + +func localSessionAdmissionKey[M MessageType](store SessionEventStore[M], sessionID string) string { + v := reflect.ValueOf(store) + if !v.IsValid() { + return "" + } + switch v.Kind() { + case reflect.Chan, reflect.Func, reflect.Map, reflect.Ptr, reflect.Slice: + if v.IsNil() { + return "" + } + return fmt.Sprintf("%T:%x/%s", store, v.Pointer(), sessionID) + default: + return fmt.Sprintf("%T:%v/%s", store, store, sessionID) + } +} + +func releaseLocalSession(key string) { + localSessionAdmission.mu.Lock() + delete(localSessionAdmission.locked, key) + localSessionAdmission.mu.Unlock() +} + +type localSessionHandle[M MessageType] struct { + key string + store SessionEventStore[M] + sessionID string + + mu sync.Mutex + closed bool +} + +func (h *localSessionHandle[M]) loadEvents(ctx context.Context, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[M], error) { + if req == nil { + req = &LoadSessionEventsRequest{} + } + return h.store.LoadEvents(ctx, h.sessionID, req) +} + +func (h *localSessionHandle[M]) appendEvents(ctx context.Context, events []*SessionEvent[M]) error { + h.mu.Lock() + if h.closed { + h.mu.Unlock() + return ErrSessionBusy + } + h.mu.Unlock() + + if err := h.store.AppendEvents(ctx, h.sessionID, events); err != nil { + return err + } + return nil +} + +func (h *localSessionHandle[M]) close(context.Context) error { + h.mu.Lock() + if h.closed { + h.mu.Unlock() + return nil + } + h.closed = true + h.mu.Unlock() + releaseLocalSession(h.key) + return nil +} diff --git a/adk/session_test.go b/adk/session_test.go new file mode 100644 index 000000000..3c6587c8e --- /dev/null +++ b/adk/session_test.go @@ -0,0 +1,5872 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// sessionHelperStore is a single-session in-memory typed session store for unit tests. +// Mirrors the EventID-based cursor semantics of session.InMemoryStore so the +// in-package tests exercise the same protocol contract. +type sessionHelperStore struct { + mu sync.Mutex + checkpoints map[string][]byte + + events []storedSessionEvent + eventIDs []string + eventIDIdx map[string]int + appendBatches [][]SessionEventKind + loadErr error + appendErr error + userMsgErr error + kindErr map[SessionEventKind]error + deleteErr error +} + +type storedSessionEvent struct { + EventID string + Kind SessionEventKind + Data []byte +} + +type blockingAppendStore struct { + sessionHelperStore + appendStarted chan struct{} + releaseAppend chan struct{} + startOnce sync.Once +} + +type publicSessionHelperStore struct { + *sessionHelperStore +} + +func (s *publicSessionHelperStore) LoadEvents(ctx context.Context, sessionID string, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + res, err := s.sessionHelperStore.LoadEventsForSession(ctx, sessionID, req) + if err != nil { + return nil, err + } + return res, nil +} + +func (s *publicSessionHelperStore) AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + return s.sessionHelperStore.AppendEventsForSession(ctx, sessionID, events) +} + +func newBlockingAppendStore() *blockingAppendStore { + return &blockingAppendStore{ + sessionHelperStore: *newSessionHelperStore(), + appendStarted: make(chan struct{}), + releaseAppend: make(chan struct{}), + } +} + +func (s *blockingAppendStore) AppendEventsForSession(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + s.startOnce.Do(func() { + close(s.appendStarted) + }) + select { + case <-s.releaseAppend: + case <-ctx.Done(): + return ctx.Err() + } + return s.sessionHelperStore.AppendEventsForSession(ctx, sessionID, events) +} + +func (s *blockingAppendStore) AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, sessionID, events) +} + +func (s *blockingAppendStore) openSession(_ context.Context, req *openSessionRequest) (*openSessionResult[*schema.Message], error) { + sessionID := "" + if req != nil { + sessionID = req.sessionID + } + return &openSessionResult[*schema.Message]{handle: &legacyMessageTestHandle{store: s, sessionID: sessionID}}, nil +} + +func (s *blockingAppendStore) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, "", events) +} + +// withTestEventID assigns a fresh UUIDv4 to the SessionEvent if its EventID is +// empty. Tests that construct SessionEvent literals directly bypass the Runner +// allocation paths, so they must still satisfy the AppendEvents wire contract. +func withTestEventID[M MessageType](se *SessionEvent[M]) *SessionEvent[M] { + if se != nil && se.EventID == "" { + se.EventID = uuid.NewString() + } + return se +} + +func withTestCommittedIdle[M MessageType](eventID string) *SessionEvent[M] { + return withTestEventID(&SessionEvent[M]{ + EventID: eventID, + Kind: SessionEventSessionStatusIdle, + Lifecycle: &LifecycleEvent{ + State: SessionRunStateIdle, + StopReason: &StopReason{Type: "end_turn"}, + }, + }) +} + +func testSequentialEventIDGenerator(prefix string) SessionEventIDGenerator[*schema.Message] { + var n int64 + return func(_ context.Context, _ *SessionEvent[*schema.Message]) (string, error) { + return fmt.Sprintf("%s%d", prefix, atomic.AddInt64(&n, 1)), nil + } +} + +// validTestPayload returns a storedSessionEvent that satisfies the AppendEvents +// wire contract (non-empty EventID) for persister-level tests that don't +// care about the SessionEvent body. +func validTestPayload() *SessionEvent[*schema.Message] { + return &SessionEvent[*schema.Message]{EventID: uuid.NewString(), Kind: SessionEventMessage, Message: schema.UserMessage("test")} +} + +func decodeStoredSessionEvents(t *testing.T, raw []storedSessionEvent) []*SessionEvent[*schema.Message] { + t.Helper() + out := make([]*SessionEvent[*schema.Message], 0, len(raw)) + for _, ep := range raw { + se, err := decodeSessionEvent[*schema.Message](ep.Data) + require.NoError(t, err) + out = append(out, se) + } + return out +} + +func filterStoredSessionEvents(t *testing.T, raw []storedSessionEvent, pred func(*SessionEvent[*schema.Message]) bool) []*SessionEvent[*schema.Message] { + t.Helper() + var out []*SessionEvent[*schema.Message] + for _, se := range decodeStoredSessionEvents(t, raw) { + if pred(se) { + out = append(out, se) + } + } + return out +} + +type testSessionAppendStore interface { + AppendEventsForSession(context.Context, string, []*SessionEvent[*schema.Message]) error +} + +func appendTestSessionEvent(t *testing.T, ctx context.Context, store testSessionAppendStore, sid string, se *SessionEvent[*schema.Message]) *SessionEvent[*schema.Message] { + t.Helper() + se = withTestEventID(se) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + return se +} + +func testMessageWithID(content string, role schema.RoleType) *schema.Message { + var msg *schema.Message + switch role { + case schema.Assistant: + msg = schema.AssistantMessage(content, nil) + default: + msg = schema.UserMessage(content) + } + EnsureMessageID(msg) + return msg +} + +func appendCommittedTestTurn(t *testing.T, ctx context.Context, store testSessionAppendStore, sid string, boundaryEventID string, contents ...string) *SessionEvent[*schema.Message] { + t.Helper() + for i, content := range contents { + role := schema.User + if i%2 == 1 { + role = schema.Assistant + } + appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Message: testMessageWithID(content, role), + }) + } + return appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + EventID: boundaryEventID, + Kind: SessionEventSessionStatusIdle, + Lifecycle: &LifecycleEvent{ + State: SessionRunStateIdle, + StopReason: &StopReason{Type: "end_turn"}, + }, + }) +} + +type runnerSessionAgent struct { + name string + inputs [][]*schema.Message + values []map[string]any + turnEnd *testTurnState[*schema.Message] +} + +type testTurnState[M MessageType] struct { + Messages []M + ToolInfos []*schema.ToolInfo + DeferredToolInfos []*schema.ToolInfo + SessionValues map[string]any +} + +func (a *runnerSessionAgent) Name(_ context.Context) string { return a.name } +func (a *runnerSessionAgent) Description(_ context.Context) string { return "runner session agent" } +func (a *runnerSessionAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + a.inputs = append(a.inputs, append([]*schema.Message{}, input.Messages...)) + a.values = append(a.values, GetSessionValues(ctx)) + go func() { + defer gen.Close() + gen.Send(&AgentEvent{ + AgentName: a.name, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: schema.AssistantMessage("ok", nil), Role: schema.Assistant}, + }, + }) + }() + return iter +} + +type streamingSessionAgent struct { + release chan struct{} + variant *SessionEventVariant[*schema.Message] +} + +func (a *streamingSessionAgent) Name(_ context.Context) string { return "streaming-session-agent" } +func (a *streamingSessionAgent) Description(_ context.Context) string { + return "streaming session agent" +} +func (a *streamingSessionAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + sr, sw := schema.Pipe[*schema.Message](1) + go func() { + defer gen.Close() + if closed := sw.Send(schema.AssistantMessage("partial", nil), nil); closed { + return + } + gen.Send(&AgentEvent{ + AgentName: a.Name(context.Background()), + Output: &AgentOutput{ + MessageOutput: &MessageVariant{IsStreaming: true, MessageStream: sr, Role: schema.Assistant}, + }, + SessionEventVariant: a.variant, + }) + <-a.release + sw.Close() + }() + return iter +} + +type erroredStreamingInterruptAgent struct { + streamErr error +} + +func (a *erroredStreamingInterruptAgent) Name(_ context.Context) string { + return "errored-streaming-interrupt-agent" +} + +func (a *erroredStreamingInterruptAgent) Description(_ context.Context) string { + return "errored streaming interrupt agent" +} + +func (a *erroredStreamingInterruptAgent) Run(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + sr, sw := schema.Pipe[*schema.Message](2) + streamErr := a.streamErr + if streamErr == nil { + streamErr = errors.New("stream failed") + } + go func() { + defer gen.Close() + sw.Send(schema.AssistantMessage("partial", nil), nil) + sw.Send(nil, streamErr) + sw.Close() + gen.Send(&AgentEvent{ + AgentName: a.Name(ctx), + Output: &AgentOutput{ + MessageOutput: &MessageVariant{IsStreaming: true, MessageStream: sr, Role: schema.Assistant}, + }, + }) + gen.Send(Interrupt(ctx, "checkpoint after errored stream")) + }() + return iter +} + +func newSessionHelperStore() *sessionHelperStore { + return &sessionHelperStore{ + checkpoints: make(map[string][]byte), + eventIDIdx: make(map[string]int), + } +} + +func (s *sessionHelperStore) Set(_ context.Context, key string, value []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + s.checkpoints[key] = append([]byte{}, value...) + return nil +} + +func (s *sessionHelperStore) Get(_ context.Context, key string) ([]byte, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.checkpoints[key] + return append([]byte{}, v...), ok, nil +} + +func (s *sessionHelperStore) Delete(_ context.Context, key string) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.deleteErr != nil { + return s.deleteErr + } + delete(s.checkpoints, key) + return nil +} + +func (s *sessionHelperStore) AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, sessionID, events) +} + +func (s *sessionHelperStore) AppendEventsForSession(_ context.Context, _ string, events []*SessionEvent[*schema.Message]) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.appendErr != nil { + return s.appendErr + } + batch := make([]SessionEventKind, 0, len(events)) + for _, e := range events { + if e == nil || e.EventID == "" { + return ErrInvalidEventID + } + if err := NormalizeSessionEventKind(e); err != nil { + return err + } + if err := s.kindErr[e.Kind]; err != nil { + return err + } + if s.userMsgErr != nil && e.Message != nil && e.Message.Role == schema.User { + return s.userMsgErr + } + if _, dup := s.eventIDIdx[e.EventID]; dup { + continue + } + batch = append(batch, e.Kind) + data, err := encodeSessionEvent(e) + if err != nil { + return err + } + s.events = append(s.events, storedSessionEvent{ + EventID: e.EventID, + Kind: e.Kind, + Data: append([]byte{}, data...), + }) + s.eventIDs = append(s.eventIDs, e.EventID) + s.eventIDIdx[e.EventID] = len(s.events) - 1 + } + if len(batch) > 0 { + s.appendBatches = append(s.appendBatches, batch) + } + return nil +} + +func (s *sessionHelperStore) LoadEvents(ctx context.Context, sessionID string, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + return s.LoadEventsForSession(ctx, sessionID, req) +} + +func (s *sessionHelperStore) LoadEventsForSession(_ context.Context, _ string, opts *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.loadErr != nil { + return nil, s.loadErr + } + if opts == nil { + opts = &LoadSessionEventsRequest{} + } + all := s.events + + if opts.Reverse { + end := len(all) + if opts.After != "" { + pos, ok := s.eventIDIdx[opts.After] + if !ok { + return nil, ErrEventIDOutOfRange + } + end = pos + } + if end <= 0 { + return &LoadSessionEventsResult[*schema.Message]{}, nil + } + kindSet := buildTestKindSet(opts.Kinds) + var out []*SessionEvent[*schema.Message] + hasMore := false + for i := end - 1; i >= 0; i-- { + if kindSet != nil { + if _, ok := kindSet[all[i].Kind]; !ok { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + event, err := decodeSessionEvent[*schema.Message](all[i].Data) + if err != nil { + return nil, err + } + out = append(out, event) + } + var next string + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &LoadSessionEventsResult[*schema.Message]{Events: out, Next: next}, nil + } + + start := 0 + if opts.After != "" { + pos, ok := s.eventIDIdx[opts.After] + if !ok { + return nil, ErrEventIDOutOfRange + } + start = pos + 1 + } + if start > len(all) { + start = len(all) + } + kindSet := buildTestKindSet(opts.Kinds) + var out []*SessionEvent[*schema.Message] + hasMore := false + for i := start; i < len(all); i++ { + if kindSet != nil { + if _, ok := kindSet[all[i].Kind]; !ok { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + event, err := decodeSessionEvent[*schema.Message](all[i].Data) + if err != nil { + return nil, err + } + out = append(out, event) + } + var next string + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &LoadSessionEventsResult[*schema.Message]{Events: out, Next: next}, nil +} + +func (s *sessionHelperStore) openSession(_ context.Context, req *openSessionRequest) (*openSessionResult[*schema.Message], error) { + sessionID := "" + if req != nil { + sessionID = req.sessionID + } + return &openSessionResult[*schema.Message]{ + handle: &testSessionHandle{store: s, sessionID: sessionID}, + }, nil +} + +func (s *sessionHelperStore) loadEvents(ctx context.Context, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + return s.LoadEventsForSession(ctx, "", req) +} + +func (s *sessionHelperStore) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, "", events) +} + +func (s *sessionHelperStore) close(context.Context) error { return nil } + +type testSessionHandle struct { + store *sessionHelperStore + sessionID string +} + +func (h *testSessionHandle) loadEvents(ctx context.Context, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + if req == nil { + req = &LoadSessionEventsRequest{} + } + return h.store.LoadEventsForSession(ctx, h.sessionID, req) +} + +func (h *testSessionHandle) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return h.store.AppendEventsForSession(ctx, h.sessionID, events) +} + +func (h *testSessionHandle) close(context.Context) error { return nil } + +type legacyMessageTestStore interface { + AppendEventsForSession(context.Context, string, []*SessionEvent[*schema.Message]) error + LoadEventsForSession(context.Context, string, *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) +} + +type legacyMessageTestHandle struct { + store legacyMessageTestStore + sessionID string +} + +func (h *legacyMessageTestHandle) loadEvents(ctx context.Context, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + if req == nil { + req = &LoadSessionEventsRequest{} + } + return h.store.LoadEventsForSession(ctx, h.sessionID, req) +} + +func (h *legacyMessageTestHandle) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return h.store.AppendEventsForSession(ctx, h.sessionID, events) +} + +func (h *legacyMessageTestHandle) close(context.Context) error { return nil } + +func mustOpenTestSession[M MessageType](t testing.TB, ctx context.Context, store SessionEventStore[M], sessionID string) sessionHandle[M] { + t.Helper() + res, err := openLocalSession(ctx, store, &openSessionRequest{sessionID: sessionID}) + require.NoError(t, err) + require.NotNil(t, res) + require.NotNil(t, res.handle) + t.Cleanup(func() { _ = res.handle.close(ctx) }) + return res.handle +} + +func buildTestKindSet(kinds []SessionEventKind) map[SessionEventKind]struct{} { + if len(kinds) == 0 { + return nil + } + set := make(map[SessionEventKind]struct{}, len(kinds)) + for _, kind := range kinds { + set[kind] = struct{}{} + } + return set +} + +func TestRunnerSessionModePrependsCommittedMessagesOnce(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sessionID := "runner-session" + firstAgent := &runnerSessionAgent{ + name: "runner-session-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("first"), schema.AssistantMessage("answer1", nil)}, + SessionValues: map[string]any{"k": "restored"}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: firstAgent, + SessionID: sessionID, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "first")) + + secondAgent := &runnerSessionAgent{ + name: "runner-session-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("first"), schema.AssistantMessage("answer1", nil), schema.UserMessage("second"), schema.AssistantMessage("answer2", nil)}, + SessionValues: map[string]any{"k": "next"}, + }, + } + runner = NewRunner(ctx, RunnerConfig{ + Agent: secondAgent, + SessionID: sessionID, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "second", WithSessionValues(map[string]any{"override": "value"}))) + + require.Len(t, secondAgent.inputs, 1) + require.Len(t, secondAgent.inputs[0], 3) + assert.Equal(t, "first", secondAgent.inputs[0][0].Content) + assert.Equal(t, "ok", secondAgent.inputs[0][1].Content) + assert.Equal(t, "second", secondAgent.inputs[0][2].Content) + require.Len(t, secondAgent.values, 1) + assert.Nil(t, secondAgent.values[0]["k"]) + assert.Equal(t, "value", secondAgent.values[0]["override"]) +} + +func TestRunnerSessionModeSkipsDuplicateEmptyModelContext(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sessionID := "runner-model-context-session" + model := &leadingSystemTestModel[*schema.Message]{response: schema.AssistantMessage("ok", nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "runner-model-context-agent", + Description: "runner model context agent", + Instruction: "You are a helpful assistant.", + Model: model, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sessionID, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "first")) + drainSessionEvents(t, runner.Query(ctx, "second")) + + result, err := store.LoadEventsForSession(ctx, sessionID, &LoadSessionEventsRequest{ + Kinds: []SessionEventKind{SessionEventModelContext}, + }) + require.NoError(t, err) + require.Len(t, result.Events, 0) +} + +func TestRunnerSessionModeSkipsDuplicateToolModelContext(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sessionID := "runner-tool-model-context-session" + model := &sessionToolCallingModel{response: schema.AssistantMessage("ok", nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "runner-tool-model-context-agent", + Description: "runner tool model context agent", + Instruction: "You are a helpful assistant.", + Model: model, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{modelContextExtraTool{}}, + }, + }, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sessionID, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "first")) + drainSessionEvents(t, runner.Query(ctx, "second")) + + result, err := store.LoadEventsForSession(ctx, sessionID, &LoadSessionEventsRequest{ + Kinds: []SessionEventKind{SessionEventModelContext}, + }) + require.NoError(t, err) + require.Len(t, result.Events, 0) +} + +func TestAttack_SessionEventIDGeneratorCoversRunnerEvents(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + prefix := "attack-runner-" + agent := &runnerSessionAgent{name: "runner-event-id-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "runner-event-id-session", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: testSequentialEventIDGenerator(prefix), + }, + }) + + drainSessionEvents(t, runner.Query(ctx, "use configured ids")) + + events := decodeStoredSessionEvents(t, store.events) + require.NotEmpty(t, events) + for _, event := range events { + require.NotEmpty(t, event.EventID) + assert.Truef(t, strings.HasPrefix(event.EventID, prefix), "event %s used unexpected ID %q", event.Kind, event.EventID) + } +} + +func TestAttack_RunnerHandlesSessionEventWithoutSessionStore(t *testing.T) { + ctx := context.Background() + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerSessionAgent{name: "runner-session-event-no-service-agent"}, + }) + + iter := runner.Query(ctx, "no managed session") + var outputs []string + var errs []error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + errs = append(errs, event.Err) + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil { + outputs = append(outputs, event.Output.MessageOutput.Message.Content) + } + } + + require.Empty(t, errs, "session envelopes emitted outside managed-session mode must not panic or surface errors") + assert.Equal(t, []string{"ok"}, outputs) +} + +// TestSessionEventIDGenerator_UserMessageBusinessID 验证：generator 可以在 +// 用户输入 message 草稿上识别业务身份并返回业务 ID（§8 UserMessage 验收）。 +func TestSessionEventIDGenerator_UserMessageBusinessID(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + const businessID = "user-order-id" + gen := func(_ context.Context, e *SessionEvent[*schema.Message]) (string, error) { + if e != nil && e.Kind == SessionEventMessage && e.Message != nil && e.Message.Role == schema.User { + return businessID, nil + } + return DefaultSessionEventIDGenerator[*schema.Message](ctx, e) + } + agent := &runnerSessionAgent{name: "user-msg-business-id-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "user-msg-business-id-session", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: gen, + }, + }) + + drainSessionEvents(t, runner.Query(ctx, "hello")) + + userMsgs := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage && se.Message != nil && se.Message.Role == schema.User + }) + require.Len(t, userMsgs, 1) + assert.Equal(t, businessID, userMsgs[0].EventID, "user input message must carry the generator-supplied business ID") +} + +func TestSessionEventIDGenerator_OutputMessageDraftBusinessID(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + const businessID = "assistant-result-id" + gen := func(_ context.Context, e *SessionEvent[*schema.Message]) (string, error) { + if e != nil && e.Kind == SessionEventMessage && e.Message != nil && e.Message.Role == schema.Assistant && e.Message.Content == "ok" { + return businessID, nil + } + return DefaultSessionEventIDGenerator[*schema.Message](ctx, e) + } + agent := &runnerSessionAgent{name: "assistant-msg-business-id-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "assistant-msg-business-id-session", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: gen, + }, + }) + + drainSessionEvents(t, runner.Query(ctx, "hello")) + + assistantMsgs := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage && se.Message != nil && se.Message.Role == schema.Assistant && se.Message.Content == "ok" + }) + require.Len(t, assistantMsgs, 1) + assert.Equal(t, businessID, assistantMsgs[0].EventID, "output message generator must see the materialized message draft") +} + +// TestSessionEventIDGenerator_ControlEventsDefaultFallthrough 验证：generator +// 仅匹配业务事件时，控制事件（status_running/status_idle 等）应通过 default +// fallthrough 拿到 UUID，而非业务 ID。 +func TestSessionEventIDGenerator_ControlEventsDefaultFallthrough(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + const businessID = "selective-user-id" + gen := func(_ context.Context, e *SessionEvent[*schema.Message]) (string, error) { + if e != nil && e.Kind == SessionEventMessage && e.Message != nil && e.Message.Role == schema.User { + return businessID, nil + } + return DefaultSessionEventIDGenerator[*schema.Message](ctx, e) + } + agent := &runnerSessionAgent{name: "control-fallthrough-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "control-fallthrough-session", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: gen, + }, + }) + + drainSessionEvents(t, runner.Query(ctx, "hi")) + + controlKinds := map[SessionEventKind]struct{}{ + SessionEventSessionStatusRunning: {}, + SessionEventSessionStatusIdle: {}, + } + controlEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + _, ok := controlKinds[se.Kind] + return ok + }) + require.NotEmpty(t, controlEvents, "expected control events (status_running / status_idle) in store") + for _, se := range controlEvents { + require.NotEmpty(t, se.EventID) + assert.NotEqual(t, businessID, se.EventID, + "control event %s must default to UUID, not adopt the user-input business ID", se.Kind) + // UUID v4 string length is 36; business ID is shorter and easily told apart. + assert.Lenf(t, se.EventID, 36, "control event %s should be a UUID (got %q)", se.Kind, se.EventID) + } +} + +// TestSessionEventIDGenerator_FailClosedOnEmpty 验证：generator 返回空 ID 时 +// runner fail closed —— 抛出 ErrSessionEventIDGeneratorEmpty 且对应草稿 event +// 不会落盘。 +func TestSessionEventIDGenerator_FailClosedOnEmpty(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + gen := func(_ context.Context, e *SessionEvent[*schema.Message]) (string, error) { + if e != nil && e.Kind == SessionEventMessage && e.Message != nil && e.Message.Role == schema.User { + return "", nil + } + return DefaultSessionEventIDGenerator[*schema.Message](ctx, e) + } + agent := &runnerSessionAgent{name: "fail-closed-empty-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "fail-closed-empty-session", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: gen, + }, + }) + + iter := runner.Query(ctx, "trigger") + var errs []error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + errs = append(errs, ev.Err) + } + } + require.NotEmpty(t, errs, "expected at least one error event from fail-closed turn") + var sawSentinel bool + for _, err := range errs { + if errors.Is(err, ErrSessionEventIDGeneratorEmpty) { + sawSentinel = true + break + } + } + require.True(t, sawSentinel, "expected ErrSessionEventIDGeneratorEmpty in error stream, got %v", errs) + + // Fail-closed: the offending user-input message must NOT be persisted. + userMsgs := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage && se.Message != nil && se.Message.Role == schema.User + }) + assert.Empty(t, userMsgs, "user input message must not be persisted when its ID allocation failed") +} + +// TestSessionEventIDGenerator_FailClosedOnError 验证：generator 返回 error 时 +// runner 同样 fail closed，错误被包装并向上抛出。 +func TestSessionEventIDGenerator_FailClosedOnError(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + genErr := errors.New("custom generator failure") + gen := func(_ context.Context, e *SessionEvent[*schema.Message]) (string, error) { + if e != nil && e.Kind == SessionEventMessage && e.Message != nil && e.Message.Role == schema.User { + return "", genErr + } + return DefaultSessionEventIDGenerator[*schema.Message](ctx, e) + } + agent := &runnerSessionAgent{name: "fail-closed-err-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "fail-closed-err-session", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: gen, + }, + }) + + iter := runner.Query(ctx, "trigger") + var errs []error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + errs = append(errs, ev.Err) + } + } + require.NotEmpty(t, errs, "expected error event when generator returns error") + var sawWrapped bool + for _, err := range errs { + if errors.Is(err, genErr) { + sawWrapped = true + break + } + } + require.True(t, sawWrapped, "expected generator error to propagate via errors.Is, got %v", errs) + + userMsgs := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage && se.Message != nil && se.Message.Role == schema.User + }) + assert.Empty(t, userMsgs, "user input message must not be persisted on generator error") +} + +func TestRunnerSessionModeRejectsPendingCheckpoint(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sessionID := "runner-pending-session" + cpBytes, err := encodeRunnerSessionCheckpoint(&runnerSessionCheckpoint{Payload: []byte("opaque")}) + require.NoError(t, err) + require.NoError(t, store.Set(ctx, sessionRunnerCheckpointID(sessionID), cpBytes)) + + agent := &runnerSessionAgent{name: "runner-session-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sessionID, + SessionStore: store, + CheckPointStore: store, + }) + iter := runner.Query(ctx, "new input") + // Run should succeed — pending checkpoint is auto-abandoned. + var sawErr bool + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + sawErr = true + } + } + require.False(t, sawErr, "Run should not return any error when pending checkpoint exists") + + // Verify agent received the input messages (no prior history to reconstruct). + require.Len(t, agent.inputs, 1) + require.Len(t, agent.inputs[0], 1) + assert.Equal(t, "new input", agent.inputs[0][0].Content) +} + +func TestAttack_RunClosesSessionHandleWhenCheckpointDecodeFails(t *testing.T) { + ctx := context.Background() + store := &publicSessionHelperStore{sessionHelperStore: newSessionHelperStore()} + service := store + sessionID := "checkpoint-decode-failure-closes-handle" + cpKey := sessionRunnerCheckpointID(sessionID) + require.NoError(t, store.Set(ctx, cpKey, []byte("not a runner checkpoint"))) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerSessionAgent{name: "checkpoint-decode-fail-agent"}, + SessionID: sessionID, + SessionStore: service, + CheckPointStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + SessionAcquireTimeout: time.Millisecond, + }, + }) + iter := runner.Query(ctx, "first") + var firstErrs []error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + firstErrs = append(firstErrs, ev.Err) + } + } + require.NotEmpty(t, firstErrs) + assert.ErrorContains(t, firstErrs[0], "failed to decode session checkpoint") + + require.NoError(t, store.Delete(ctx, cpKey)) + iter = runner.Query(ctx, "second") + var secondErrs []error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + secondErrs = append(secondErrs, ev.Err) + } + } + require.Empty(t, secondErrs, "session handle must be released after checkpoint decode failure") +} + +func TestRunnerSessionModeDeleteCheckpointFailureIsReported(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + persister := newSessionEventPersister[*schema.Message]( + ctx, + store, + "delete-fail-session", + ) + checkPointID := "delete-fail-checkpoint" + store.deleteErr = errors.New("delete failed") + + res := &sessionTurnResult[*schema.Message]{ + persister: persister, + sessionState: &runnerSessionRunState[*schema.Message]{ + enabled: true, + sessionID: "delete-fail-session", + sessionStore: store, + }, + store: store, + checkPointID: &checkPointID, + } + + err := res.finalize(ctx) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to delete session checkpoint") +} + +func TestModelContextEvent_JSONLikeRoundTrip(t *testing.T) { + modelCtx := &ModelContextEvent{ + ToolInfos: []*schema.ToolInfo{ + { + Name: "lookup", + Desc: "lookup tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{"q": {Type: schema.String}}), + }, + }, + } + + se := &SessionEvent[*schema.Message]{Kind: SessionEventModelContext, ModelContext: modelCtx} + data, err := encodeSessionEvent(withTestEventID(se)) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.ModelContext) + require.Len(t, decoded.ModelContext.ToolInfos, 1) + assert.Equal(t, "lookup", decoded.ModelContext.ToolInfos[0].Name) +} + +func TestRunnerSessionStreamingDoesNotBlockLiveEvent(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + agent := &streamingSessionAgent{release: make(chan struct{})} + release := func() { + select { + case <-agent.release: + default: + close(agent.release) + } + } + defer release() + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: "streaming-session", + SessionStore: store, + }) + + iter := runner.Query(ctx, "start") + type nextResult struct { + event *AgentEvent + ok bool + } + nextCh := make(chan nextResult, 1) + go func() { + event, ok := iter.Next() + nextCh <- nextResult{event: event, ok: ok} + }() + + var res nextResult + select { + case res = <-nextCh: + case <-time.After(200 * time.Millisecond): + t.Fatal("managed session persistence blocked live streaming event delivery") + } + + require.True(t, res.ok) + require.NoError(t, res.event.Err) + require.NotNil(t, res.event.Output) + require.NotNil(t, res.event.Output.MessageOutput) + require.True(t, res.event.Output.MessageOutput.IsStreaming) + require.NotNil(t, res.event.Output.MessageOutput.MessageStream) + + msg, err := res.event.Output.MessageOutput.MessageStream.Recv() + require.NoError(t, err) + assert.Equal(t, "partial", msg.Content) + + release() + drainSessionEvents(t, iter) +} + +func TestRunnerSessionStreamingRefAllocatesMissingEventID(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + agent := &streamingSessionAgent{ + release: make(chan struct{}), + variant: &SessionEventVariant[*schema.Message]{ + MessageStreamRef: &MessageStreamRef{Kind: SessionEventMessage}, + }, + } + release := func() { + select { + case <-agent.release: + default: + close(agent.release) + } + } + defer release() + + const businessID = "stream-business-id" + var sawStreamDraft bool + gen := func(ctx context.Context, e *SessionEvent[*schema.Message]) (string, error) { + if e != nil && e.Kind == SessionEventMessage && e.Message == nil { + sawStreamDraft = true + assert.False(t, e.Timestamp.IsZero()) + return businessID, nil + } + return DefaultSessionEventIDGenerator[*schema.Message](ctx, e) + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: "streaming-ref-session", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: gen, + }, + }) + + iter := runner.Query(ctx, "start", WithTimelineEvents()) + var event *AgentEvent + for { + ev, ok := iter.Next() + require.True(t, ok) + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.IsStreaming { + event = ev + break + } + } + require.NotNil(t, event.SessionEventVariant) + ref := event.SessionEventVariant.MessageStreamRef + require.NotNil(t, ref) + assert.Equal(t, businessID, ref.EventID) + assert.Equal(t, SessionEventMessage, ref.Kind) + assert.False(t, ref.Timestamp.IsZero()) + + msg, err := event.Output.MessageOutput.MessageStream.Recv() + require.NoError(t, err) + assert.Equal(t, "partial", msg.Content) + release() + drainSessionEvents(t, iter) + + require.True(t, sawStreamDraft) + messages := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage && se.Message != nil && se.Message.Content == "partial" + }) + require.Len(t, messages, 1) + assert.Equal(t, businessID, messages[0].EventID) + assert.Equal(t, ref.Timestamp, messages[0].Timestamp) +} + +func TestRunnerSessionPersistsIncompleteStreamingMessageBeforeCheckpoint(t *testing.T) { + ctx := context.Background() + + tests := []struct { + name string + streamErr error + }{ + {name: "stream canceled", streamErr: ErrStreamCanceled}, + {name: "will retry", streamErr: &WillRetryError{ErrStr: "retry", RetryAttempt: 1}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + store := newSessionHelperStore() + agent := &erroredStreamingInterruptAgent{streamErr: tt.streamErr} + checkpointID := "errored-stream-checkpoint-" + strings.ReplaceAll(tt.name, " ", "-") + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: "errored-stream-session-" + tt.name, + SessionStore: store, + CheckPointStore: store, + }) + + iter := runner.Query(ctx, "start", WithCheckPointID(checkpointID)) + var sawInterrupt bool + var sawStreamErr bool + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming { + for { + _, err := event.Output.MessageOutput.MessageStream.Recv() + if err == nil { + continue + } + require.NotEqual(t, io.EOF, err) + sawStreamErr = true + break + } + } + if event.Action != nil && event.Action.Interrupted != nil { + sawInterrupt = true + } + } + + require.True(t, sawStreamErr) + require.True(t, sawInterrupt) + _, exists, err := store.Get(ctx, checkpointID) + require.NoError(t, err) + require.True(t, exists, "checkpoint should still be saved after incomplete message stream") + + persistedPartialMessages := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage && + se.Message != nil && + se.Message.Role == schema.Assistant && + se.Message.Content == "partial" + }) + assert.Empty(t, persistedPartialMessages) + incompleteMessages := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessageStreamIncomplete + }) + require.Len(t, incompleteMessages, 1) + require.NotNil(t, incompleteMessages[0].MessageStreamIncomplete) + assert.Equal(t, "partial", incompleteMessages[0].MessageStreamIncomplete.Message.Content) + assert.Contains(t, incompleteMessages[0].MessageStreamIncomplete.Error, tt.streamErr.Error()) + }) + } +} + +func drainSessionEvents(t *testing.T, iter *AsyncIterator[*AgentEvent]) { + t.Helper() + for { + event, ok := iter.Next() + if !ok { + return + } + require.NoError(t, event.Err) + } +} + +// runnerInterruptAgent: produces an interrupt on first Run; emits "resumed ok" on Resume. +type runnerInterruptAgent struct { + callCount int32 +} + +func (a *runnerInterruptAgent) Name(_ context.Context) string { return "InterruptAgent" } +func (a *runnerInterruptAgent) Description(_ context.Context) string { return "runner interrupt agent" } + +func (a *runnerInterruptAgent) Run(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + atomic.AddInt32(&a.callCount, 1) + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + event := Interrupt(ctx, "confirm?") + gen.Send(event) + }() + return iter +} + +func (a *runnerInterruptAgent) Resume(ctx context.Context, info *ResumeInfo, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + atomic.AddInt32(&a.callCount, 1) + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(&AgentEvent{ + AgentName: "InterruptAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage("resumed ok", nil), + Role: schema.Assistant, + }, + }, + }) + }() + return iter +} + +type runnerCheckpointSanitizeAgent struct{} + +func (a *runnerCheckpointSanitizeAgent) Name(_ context.Context) string { + return "CheckpointSanitizeAgent" +} + +func (a *runnerCheckpointSanitizeAgent) Description(_ context.Context) string { + return "session checkpoint sanitizer test agent" +} + +func (a *runnerCheckpointSanitizeAgent) Run(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(&AgentEvent{ + AgentName: "CheckpointSanitizeAgent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: "checkpoint-session-only", + Kind: SessionEventSessionStatusRunning, + Lifecycle: &LifecycleEvent{ + State: SessionRunStateRunning, + }, + }, + }, + }) + gen.Send(&AgentEvent{ + AgentName: "CheckpointSanitizeAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage("mixed output", nil), + Role: schema.Assistant, + }, + }, + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: "checkpoint-output", + Kind: SessionEventMessage, + Message: schema.AssistantMessage("mixed output", nil), + }, + }, + }) + gen.Send(Interrupt(ctx, "confirm?")) + }() + return iter +} + +func TestRunnerSessionModeResumeWithEmptyCheckpointID(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sessionID := "resume-test" + + agent := &runnerInterruptAgent{} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sessionID, + SessionStore: store, + CheckPointStore: store, + }) + + iter := runner.Query(ctx, "hello") + var sawInterrupt bool + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + sawInterrupt = true + } + } + require.True(t, sawInterrupt) + + resumeIter, err := runner.Resume(ctx, "") + require.NoError(t, err) + var gotResumedOK bool + for { + event, ok := resumeIter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.Message != nil && + event.Output.MessageOutput.Message.Content == "resumed ok" { + gotResumedOK = true + } + } + assert.True(t, gotResumedOK) +} + +func TestRunnerSessionModeFlushFailurePreventsCommit(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + store.appendErr = errors.New("disk full") + + agent := &runnerSessionAgent{ + name: "flush-fail-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("done", nil)}, + }, + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "flush-fail-session", + SessionStore: store, + }) + + iter := runner.Query(ctx, "trigger") + var lastErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + lastErr = event.Err + } + } + + require.Error(t, lastErr) + assert.Contains(t, lastErr.Error(), "disk full") +} + +func TestRunnerSessionSyncModeBlocksDeliveryUntilAppendCompletes(t *testing.T) { + ctx := context.Background() + store := newBlockingAppendStore() + agent := &runnerSessionAgent{ + name: "sync-block-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("ok", nil)}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "sync-block-session", + SessionStore: store, + }) + + iterCh := make(chan *AsyncIterator[*AgentEvent], 1) + go func() { + iterCh <- runner.Query(ctx, "trigger") + }() + select { + case <-iterCh: + t.Fatal("query returned before pre-run control append completed") + case <-time.After(50 * time.Millisecond): + } + events := make(chan *AgentEvent, 1) + + select { + case <-store.appendStarted: + case <-time.After(500 * time.Millisecond): + t.Fatal("sync persistence did not start appending") + } + close(store.releaseAppend) + iter := <-iterCh + go func() { + ev, ok := iter.Next() + if !ok { + events <- nil + return + } + events <- ev + }() + firstEvent := <-events + var sawOutput bool + if firstEvent != nil { + require.NoError(t, firstEvent.Err) + if firstEvent.Output != nil && firstEvent.Output.MessageOutput != nil && + firstEvent.Output.MessageOutput.Message != nil && + firstEvent.Output.MessageOutput.Message.Content == "ok" { + sawOutput = true + } + } + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.Message != nil && + ev.Output.MessageOutput.Message.Content == "ok" { + sawOutput = true + break + } + } + assert.True(t, sawOutput, "expected output after sync append completed") +} + +func TestRunnerSessionSyncModeAppendFailureSuppressesOutput(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + store.appendErr = errors.New("sync append failed") + agent := &runnerSessionAgent{ + name: "sync-fail-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("ok", nil)}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "sync-fail-session", + SessionStore: store, + }) + + iter := runner.Query(ctx, "trigger") + var lastErr error + var sawOutput bool + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + lastErr = ev.Err + } + if ev.Output != nil && ev.Output.MessageOutput != nil { + sawOutput = true + } + } + + require.Error(t, lastErr) + assert.Contains(t, lastErr.Error(), "sync append failed") + assert.False(t, sawOutput, "sync mode must not deliver output after append failure") +} + +func TestSessionPersister_EnqueueAfterFlush(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + + persister := newSessionEventPersister[*schema.Message](ctx, store, "enqueue-after-flush") + + require.NoError(t, persister.closeAndWait()) + assert.NoError(t, persister.enqueueAsync(validTestPayload())) + require.NoError(t, persister.closeAndWait()) + assert.Len(t, store.events, 1) +} + +// TestSessionPersister_EmptyPayloadSkipped verifies enqueue silently discards +// records with empty payload. +func TestSessionPersister_EmptyPayloadSkipped(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + + persister := newSessionEventPersister[*schema.Message](ctx, store, "empty-payload") + + assert.NoError(t, persister.enqueueAsync(nil)) + assert.NoError(t, persister.enqueueAsync(&SessionEvent[*schema.Message]{})) + + se := makeInputSessionEvent(schema.UserMessage("real")) + se.EventID = uuid.NewString() + require.NoError(t, persister.enqueueAsync(se)) + + require.NoError(t, persister.closeAndWait()) + require.Len(t, store.events, 1, "only the real event should be persisted") +} + +func TestSessionPersister_AsyncEnqueueFlushesOnClose(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + persister := newSessionEventPersister[*schema.Message](ctx, store, "async-enqueue") + + require.NoError(t, persister.enqueueAsync(validTestPayload())) + store.mu.Lock() + assert.Len(t, store.events, 0, "async annotations stay pending until a boundary or final flush") + store.mu.Unlock() + + require.NoError(t, persister.closeAndWait()) + store.mu.Lock() + assert.Len(t, store.events, 1, "closeAndWait flushes pending annotations") + store.mu.Unlock() +} + +func TestSessionPersister_CommitBoundaryFlushesPendingBatchShape(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + persister := newSessionEventPersister[*schema.Message](ctx, store, "boundary-shape") + + annotation := withTestEventID(&SessionEvent[*schema.Message]{ + Kind: SessionEventKind(SessionEventExtensionPrefix + "annotation"), + Extension: &SessionExtensionEvent{}, + }) + message := withTestEventID(&SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Message: schema.AssistantMessage("durable", nil), + }) + + require.NoError(t, persister.enqueueAsync(annotation)) + require.NoError(t, persister.commitBoundary(message)) + require.NoError(t, persister.closeAndWait()) + + assert.Equal(t, [][]SessionEventKind{ + {annotation.Kind}, + {SessionEventMessage}, + }, store.appendBatches) +} + +func TestSessionPersister_CommitBoundaryPreservesPendingOnFlushFailure(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + persister := newSessionEventPersister[*schema.Message](ctx, store, "boundary-fail") + + annotation := withTestEventID(&SessionEvent[*schema.Message]{ + Kind: SessionEventKind(SessionEventExtensionPrefix + "annotation"), + Extension: &SessionExtensionEvent{}, + }) + message := withTestEventID(&SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Message: schema.AssistantMessage("durable", nil), + }) + require.NoError(t, persister.enqueueAsync(annotation)) + + store.appendErr = errors.New("flush failed") + err := persister.commitBoundary(message) + require.Error(t, err) + assert.Contains(t, err.Error(), "flush failed") + assert.Empty(t, store.events) + require.Len(t, persister.pending, 1) + assert.Equal(t, annotation.EventID, persister.pending[0].EventID) +} + +func TestRunnerSessionDurableBoundaryBatchShape(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + agent := &runnerSessionAgent{name: "boundary-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "boundary-session", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hello") + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + } + + assert.Equal(t, [][]SessionEventKind{ + {SessionEventSessionStatusRunning}, + {SessionEventMessage}, + {SessionEventMessage}, + {SessionEventSessionStatusIdle}, + }, store.appendBatches) +} + +func TestRunnerSessionInputMessageBoundaryFailureStopsBeforeAgent(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + store.userMsgErr = errors.New("input append failed") + agent := &runnerSessionAgent{name: "input-boundary-agent"} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "input-boundary-session", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hello") + event, ok := iter.Next() + require.True(t, ok) + require.Error(t, event.Err) + assert.Contains(t, event.Err.Error(), "input append failed") + _, ok = iter.Next() + assert.False(t, ok) + assert.Empty(t, agent.inputs, "agent must not execute when input message boundary append fails") + assert.Equal(t, [][]SessionEventKind{{SessionEventSessionStatusRunning}}, store.appendBatches) +} + +func TestSessionPersister_DirectAppendNoRetryAndLatch(t *testing.T) { + t.Run("transient failure is not retried by runner", func(t *testing.T) { + ctx := context.Background() + store := &transientFailStore{ + sessionHelperStore: *newSessionHelperStore(), + failsLeft: 2, + appendErrVal: errors.New("transient"), + } + persister := newSessionEventPersister[*schema.Message](ctx, store, "no-retry") + + err := persister.commitBoundary(validTestPayload()) + require.Error(t, err) + assert.Contains(t, err.Error(), "transient") + assert.Equal(t, 1, store.getAppendCalls()) + }) + + t.Run("permanent failure latched", func(t *testing.T) { + ctx := context.Background() + store := &transientFailStore{ + sessionHelperStore: *newSessionHelperStore(), + failsLeft: 100, + appendErrVal: errors.New("permanent"), + } + persister := newSessionEventPersister[*schema.Message](ctx, store, "latch") + + err := persister.commitBoundary(validTestPayload()) + require.Error(t, err) + assert.Contains(t, err.Error(), "permanent") + assert.Equal(t, 1, store.getAppendCalls()) + + err = persister.enqueueAsync(validTestPayload()) + require.Error(t, err) + assert.Contains(t, err.Error(), "permanent") + assert.Equal(t, 1, store.getAppendCalls(), "latched failure must prevent later appends") + assert.Error(t, persister.closeAndWait()) + }) +} + +// TestModelContextEvent_GobRoundtripNilFields verifies gob roundtrip preserves nil semantics. +func TestModelContextEvent_GobRoundtripNilFields(t *testing.T) { + se := &SessionEvent[*schema.Message]{Kind: SessionEventModelContext, ModelContext: &ModelContextEvent{}} + encoded, err := encodeSessionEvent(withTestEventID(se)) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](encoded) + require.NoError(t, err) + require.NotNil(t, decoded.ModelContext) + assert.Nil(t, decoded.ModelContext.ToolInfos) + assert.Nil(t, decoded.ModelContext.DeferredToolInfos) +} + +func TestNormalizeSessionConfig_Variations(t *testing.T) { + cfg := normalizeSessionConfig[*schema.Message](nil) + assert.NotNil(t, cfg.EventIDGenerator) + assert.Equal(t, defaultSessionAcquireTimeout, cfg.SessionAcquireTimeout) + + cfg = normalizeSessionConfig(&SessionConfig[*schema.Message]{}) + assert.NotNil(t, cfg.EventIDGenerator) + assert.Equal(t, defaultSessionAcquireTimeout, cfg.SessionAcquireTimeout) + + customGen := func(context.Context, *SessionEvent[*schema.Message]) (string, error) { + return "custom-id", nil + } + cfg = normalizeSessionConfig(&SessionConfig[*schema.Message]{ + EventIDGenerator: customGen, + SessionAcquireTimeout: 200 * time.Millisecond, + }) + assert.Equal(t, 200*time.Millisecond, cfg.SessionAcquireTimeout) + id, err := cfg.EventIDGenerator(context.Background(), nil) + require.NoError(t, err) + assert.Equal(t, "custom-id", id) +} + +type countingSerializer struct { + inner schema.Serializer + marshalCalls int32 + unmarshalCalls int32 +} + +func newCountingSerializer() *countingSerializer { + return &countingSerializer{inner: &schema.HumanReadableSerializer{}} +} + +func (s *countingSerializer) Marshal(v any) ([]byte, error) { + atomic.AddInt32(&s.marshalCalls, 1) + return s.inner.Marshal(v) +} + +func (s *countingSerializer) Unmarshal(data []byte, v any) error { + atomic.AddInt32(&s.unmarshalCalls, 1) + return s.inner.Unmarshal(data, v) +} + +func TestSessionEvent_HumanReadableSerializerDirectRoundTrip(t *testing.T) { + serializer := &schema.HumanReadableSerializer{} + se := &SessionEvent[*schema.Message]{ + EventID: "serializer-direct", + Kind: SessionEventSessionStatusIdle, + Lifecycle: &LifecycleEvent{ + State: SessionRunStateIdle, + }, + } + + data, err := serializer.Marshal(se) + require.NoError(t, err) + + var decoded SessionEvent[*schema.Message] + require.NoError(t, serializer.Unmarshal(data, &decoded)) + require.NoError(t, NormalizeSessionEventKind(&decoded)) + assert.Equal(t, se.EventID, decoded.EventID) + assert.Equal(t, se.Kind, decoded.Kind) +} + +// --- New tests covering the design doc --- + +func TestSessionEvent_HumanReadableRoundTrip(t *testing.T) { + t.Run("Message", func(t *testing.T) { + msg := schema.UserMessage("hello") + EnsureMessageID(msg) + se := &SessionEvent[*schema.Message]{Message: msg} + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.Message) + assert.Equal(t, "hello", decoded.Message.Content) + assert.Equal(t, GetMessageID(msg), GetMessageID(decoded.Message)) + }) + + t.Run("MessagesReplaced", func(t *testing.T) { + msgs := []*schema.Message{schema.UserMessage("a"), schema.AssistantMessage("b", nil)} + for _, m := range msgs { + EnsureMessageID(m) + } + se := &SessionEvent[*schema.Message]{MessagesReplaced: &msgs} + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.MessagesReplaced) + assert.Equal(t, 2, len(*decoded.MessagesReplaced)) + assert.Equal(t, "a", (*decoded.MessagesReplaced)[0].Content) + }) + + t.Run("MessageUpdated", func(t *testing.T) { + updated := schema.AssistantMessage("placeholder", nil) + EnsureMessageID(updated) + se := &SessionEvent[*schema.Message]{ + MessageUpdated: &MessageUpdatedEvent[*schema.Message]{ + MessageID: GetMessageID(updated), + Message: updated, + }, + } + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.MessageUpdated) + assert.Equal(t, GetMessageID(updated), decoded.MessageUpdated.MessageID) + assert.Equal(t, "placeholder", decoded.MessageUpdated.Message.Content) + }) + + t.Run("MessageInserted", func(t *testing.T) { + inserted := schema.UserMessage("agentsmd content") + EnsureMessageID(inserted) + se := &SessionEvent[*schema.Message]{ + MessageInserted: &MessageInsertedEvent[*schema.Message]{ + Message: inserted, + BeforeMessageID: "anchor-id", + }, + } + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.MessageInserted) + assert.Equal(t, "anchor-id", decoded.MessageInserted.BeforeMessageID) + assert.Equal(t, "agentsmd content", decoded.MessageInserted.Message.Content) + }) + + t.Run("MessagesDeleted", func(t *testing.T) { + se := &SessionEvent[*schema.Message]{ + MessagesDeleted: &MessagesDeletedEvent{MessageIDs: []string{"m1", "m2"}}, + } + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.MessagesDeleted) + assert.Equal(t, SessionEventMessagesDeleted, decoded.Kind) + assert.Equal(t, []string{"m1", "m2"}, decoded.MessagesDeleted.MessageIDs) + }) +} + +// TestApplySessionEvent verifies all variants of the event-applier. +func TestApplySessionEvent(t *testing.T) { + makeMsg := func(content string) *schema.Message { + m := schema.UserMessage(content) + EnsureMessageID(m) + return m + } + + t.Run("Message appends", func(t *testing.T) { + var msgs []*schema.Message + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{Message: makeMsg("a")}) + require.NoError(t, err) + require.Len(t, msgs, 1) + }) + + t.Run("MessagesReplaced replaces wholesale", func(t *testing.T) { + msgs := []*schema.Message{makeMsg("old")} + repl := []*schema.Message{makeMsg("new1"), makeMsg("new2")} + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{MessagesReplaced: &repl}) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, "new1", msgs[0].Content) + }) + + t.Run("MessageUpdated replaces in place", func(t *testing.T) { + target := makeMsg("orig") + msgs := []*schema.Message{makeMsg("a"), target, makeMsg("b")} + newMsg := schema.AssistantMessage("placeholder", nil) + newMsg.Extra = map[string]any{} + // Force same ID + setMessageIDForTest(newMsg, GetMessageID(target)) + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessageUpdated: &MessageUpdatedEvent[*schema.Message]{ + MessageID: GetMessageID(target), + Message: newMsg, + }, + }) + require.NoError(t, err) + assert.Equal(t, "placeholder", msgs[1].Content) + }) + + t.Run("MessageUpdated identity mismatch", func(t *testing.T) { + target := makeMsg("orig") + msgs := []*schema.Message{target} + other := makeMsg("other") + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessageUpdated: &MessageUpdatedEvent[*schema.Message]{ + MessageID: GetMessageID(target), + Message: other, // has its own different ID + }, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "identity mismatch") + }) + + t.Run("MessageInserted before anchor", func(t *testing.T) { + anchor := makeMsg("anchor") + msgs := []*schema.Message{makeMsg("a"), anchor, makeMsg("b")} + ins := makeMsg("inserted") + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessageInserted: &MessageInsertedEvent[*schema.Message]{ + Message: ins, + BeforeMessageID: GetMessageID(anchor), + }, + }) + require.NoError(t, err) + require.Len(t, msgs, 4) + assert.Equal(t, "inserted", msgs[1].Content) + assert.Equal(t, "anchor", msgs[2].Content) + }) + + t.Run("MessageInserted append at end", func(t *testing.T) { + msgs := []*schema.Message{makeMsg("a")} + ins := makeMsg("appended") + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessageInserted: &MessageInsertedEvent[*schema.Message]{Message: ins, BeforeMessageID: ""}, + }) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, "appended", msgs[1].Content) + }) + + t.Run("MessageInserted missing anchor errors", func(t *testing.T) { + msgs := []*schema.Message{makeMsg("a")} + ins := makeMsg("ghost") + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessageInserted: &MessageInsertedEvent[*schema.Message]{ + Message: ins, + BeforeMessageID: "no-such-anchor", + }, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "anchor message") + }) + + t.Run("MessageUpdated missing target errors", func(t *testing.T) { + msgs := []*schema.Message{makeMsg("a")} + other := makeMsg("other") + setMessageIDForTest(other, "ghost-id") + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessageUpdated: &MessageUpdatedEvent[*schema.Message]{ + MessageID: "ghost-id", + Message: other, + }, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "not found for update") + }) + + t.Run("MessagesDeleted removes multiple messages", func(t *testing.T) { + a := makeMsg("a") + b := makeMsg("b") + c := makeMsg("c") + d := makeMsg("d") + msgs := []*schema.Message{a, b, c, d} + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessagesDeleted: &MessagesDeletedEvent{MessageIDs: []string{GetMessageID(b), GetMessageID(d)}}, + }) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, "a", msgs[0].Content) + assert.Equal(t, "c", msgs[1].Content) + }) + + t.Run("MessagesDeleted missing target errors", func(t *testing.T) { + a := makeMsg("a") + b := makeMsg("b") + c := makeMsg("c") + msgs := []*schema.Message{a, b, c} + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessagesDeleted: &MessagesDeletedEvent{MessageIDs: []string{GetMessageID(b), "ghost-id"}}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "ghost-id") + assert.Equal(t, []*schema.Message{a, b, c}, msgs) + }) + + t.Run("MessagesDeleted rejects empty and duplicate ids", func(t *testing.T) { + msgs := []*schema.Message{makeMsg("a")} + err := applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessagesDeleted: &MessagesDeletedEvent{}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "must not be empty") + + err = applySessionEvent(&msgs, &SessionEvent[*schema.Message]{ + MessagesDeleted: &MessagesDeletedEvent{MessageIDs: []string{"dup", "dup"}}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "duplicate") + }) +} + +func setMessageIDForTest(msg *schema.Message, id string) { + if msg.Extra == nil { + msg.Extra = map[string]any{} + } + msg.Extra["_eino_msg_id"] = id +} + +// TestStripSessionEventFields verifies all session-internal fields are stripped. +func TestStripSessionEventFields(t *testing.T) { + t.Run("non-session-internal event passes through", func(t *testing.T) { + ev := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: schema.AssistantMessage("hi", nil), Role: schema.Assistant}, + }, + } + stripped := stripSessionEventFields(ev) + require.NotNil(t, stripped) + assert.Equal(t, "hi", stripped.Output.MessageOutput.Message.Content) + }) + + t.Run("SessionEvent-only event drops to nil", func(t *testing.T) { + ev := &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventModelContext, + ModelContext: &ModelContextEvent{}, + }, + }, + } + stripped := stripSessionEventFields(ev) + assert.Nil(t, stripped) + }) + + t.Run("message mutation SessionEvent-only event drops to nil", func(t *testing.T) { + msgs := []*schema.Message{schema.UserMessage("x")} + ev := &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessagesReplaced, + MessagesReplaced: &msgs, + }, + }, + } + stripped := stripSessionEventFields(ev) + assert.Nil(t, stripped) + }) + + t.Run("Err with SessionEvent keeps Err", func(t *testing.T) { + ev := &AgentEvent{ + Err: errors.New("visible"), + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + SessionID: "child-1", + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventModelContext, + ModelContext: &ModelContextEvent{}, + }, + }, + } + stripped := stripSessionEventFields(ev) + require.NotNil(t, stripped) + assert.Nil(t, stripped.SessionEventVariant) + assert.EqualError(t, stripped.Err, "visible") + }) + + t.Run("SessionEventVariant with SessionID alone is stripped", func(t *testing.T) { + ev := &AgentEvent{SessionEventVariant: &SessionEventVariant[*schema.Message]{SessionID: "child-1"}} + stripped := stripSessionEventFields(ev) + assert.Nil(t, stripped) + }) +} + +func TestSessionEventTimestamp(t *testing.T) { + ts := time.Date(2026, 5, 22, 10, 2, 0, 0, time.UTC) + msg := schema.AssistantMessage("hi", nil) + EnsureMessageID(msg) + event := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: msg, Role: schema.Assistant}, + }, + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Timestamp: ts, + Kind: SessionEventMessage, + Message: msg, + }, + }, + } + + se := toSessionEvent(event) + require.NotNil(t, se) + assert.Equal(t, ts, se.Timestamp) + + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + assert.Equal(t, ts, decoded.Timestamp) +} + +// TestReconstructFromEventLog_EmptySession verifies empty-session reconstruction. +func TestReconstructFromEventLog_EmptySession(t *testing.T) { + store := newSessionHelperStore() + ctx := context.Background() + result, err := reconstructSessionState[*schema.Message](ctx, store, "empty", defaultLoadPageSize) + require.NoError(t, err) + assert.Nil(t, result) +} + +// TestReconstructFromEventLog_MultiTurn verifies multi-turn reconstruction. +func TestReconstructFromEventLog_MultiTurn(t *testing.T) { + store := newSessionHelperStore() + ctx := context.Background() + sid := "multi-turn" + + // Turn 1: input "Q1" + output "A1" + q1 := schema.UserMessage("Q1") + EnsureMessageID(q1) + a1 := schema.AssistantMessage("A1", nil) + EnsureMessageID(a1) + for _, m := range []*schema.Message{q1, a1} { + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + // Turn 2: input "Q2" + output "A2" + q2 := schema.UserMessage("Q2") + EnsureMessageID(q2) + a2 := schema.AssistantMessage("A2", nil) + EnsureMessageID(a2) + for _, m := range []*schema.Message{q2, a2} { + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 4) + assert.Equal(t, "Q1", result.state.Messages[0].Content) + assert.Equal(t, "A1", result.state.Messages[1].Content) + assert.Equal(t, "Q2", result.state.Messages[2].Content) + assert.Equal(t, "A2", result.state.Messages[3].Content) + + // Verify pagination: use page size 2 so that 4 events require multiple pages. + result2, err := reconstructSessionState[*schema.Message](ctx, store, sid, 2) + require.NoError(t, err) + require.NotNil(t, result2) + require.NotNil(t, result2.state) + require.Len(t, result2.state.Messages, 4) + assert.Equal(t, "Q1", result2.state.Messages[0].Content) + assert.Equal(t, "A1", result2.state.Messages[1].Content) + assert.Equal(t, "Q2", result2.state.Messages[2].Content) + assert.Equal(t, "A2", result2.state.Messages[3].Content) +} + +func TestReconstructFromEventLog_CorruptEventReturnsError(t *testing.T) { + store := newSessionHelperStore() + ctx := context.Background() + sid := "corrupt-event" + + msg := schema.UserMessage("valid") + EnsureMessageID(msg) + se := withTestEventID(&SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Message: msg, + }) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + + corruptPayload := []byte(`{"event_id":"` + uuid.NewString() + `","kind":"message","message":` + "\x00\xff invalid json") + require.False(t, json.Valid(corruptPayload), "payload must be invalid JSON") + corruptID := uuid.NewString() + store.mu.Lock() + store.events = append(store.events, storedSessionEvent{EventID: corruptID, Kind: SessionEventMessage, Data: corruptPayload}) + store.eventIDs = append(store.eventIDs, corruptID) + store.eventIDIdx[corruptID] = len(store.events) - 1 + store.mu.Unlock() + + _, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.Error(t, err, "corrupt event must cause reconstruction failure") +} + +// TestReconstructFromEventLog_WithSummarizationBoundary: events before +// MessagesReplaced are ignored; reconstruction starts from boundary. +func TestReconstructFromEventLog_WithSummarizationBoundary(t *testing.T) { + store := newSessionHelperStore() + ctx := context.Background() + sid := "with-boundary" + + // Pre-boundary events (should be ignored). + for i := 0; i < 3; i++ { + m := schema.UserMessage("pre") + EnsureMessageID(m) + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + // Boundary: summary of all messages. + summary := schema.UserMessage("summary") + EnsureMessageID(summary) + repl := []*schema.Message{summary} + se := withTestEventID(&SessionEvent[*schema.Message]{MessagesReplaced: &repl}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + + // Post-boundary events. + post := schema.AssistantMessage("post", nil) + EnsureMessageID(post) + se = withTestEventID(&SessionEvent[*schema.Message]{Message: post}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 2) + assert.Equal(t, "summary", result.state.Messages[0].Content) + assert.Equal(t, "post", result.state.Messages[1].Content) +} + +func TestSessionRollbackEventRoundTrip(t *testing.T) { + se := &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Kind: SessionEventRollback, + Rollback: &SessionRollbackEvent{ + ToEventID: "turn-end-1", + PreviousHeadCommitEventID: "turn-end-2", + }, + } + data, err := encodeSessionEvent(se) + require.NoError(t, err) + + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.Rollback) + assert.Equal(t, SessionEventRollback, decoded.Kind) + assert.Equal(t, "turn-end-1", decoded.Rollback.ToEventID) + assert.Equal(t, "turn-end-2", decoded.Rollback.PreviousHeadCommitEventID) +} + +func TestAttack_RollbackSessionUsesConfiguredEventIDGenerator(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "rollback-event-id-generator" + appendCommittedTestTurn(t, ctx, store, sid, "turn-1", "Q1", "A1") + appendCommittedTestTurn(t, ctx, store, sid, "turn-2", "Q2", "A2") + + require.NoError(t, RollbackSession[*schema.Message]( + ctx, + store, + sid, + "turn-1", + WithRollbackEventIDGenerator(testSequentialEventIDGenerator("attack-rollback-")), + )) + + rollbackEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventRollback + }) + require.Len(t, rollbackEvents, 1) + assert.Equal(t, "attack-rollback-1", rollbackEvents[0].EventID) +} + +func TestRollbackSessionReconstructionHidesDeadBranchAndKeepsNewSuffix(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "rollback-reconstruct" + + t1 := appendCommittedTestTurn(t, ctx, store, sid, "turn-1", "Q1", "A1") + t2 := appendCommittedTestTurn(t, ctx, store, sid, "turn-2", "Q2", "A2") + require.NoError(t, RollbackSession[*schema.Message]( + ctx, + store, + sid, + "turn-1", + WithRollbackSessionCheckPointStore[*schema.Message](store), + WithRollbackSessionExpectedHeadEventID[*schema.Message]("turn-2"), + )) + appendCommittedTestTurn(t, ctx, store, sid, "turn-3", "Q3", "A3") + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, 2) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 4) + assert.Equal(t, "Q1", result.state.Messages[0].Content) + assert.Equal(t, "A1", result.state.Messages[1].Content) + assert.Equal(t, "Q3", result.state.Messages[2].Content) + assert.Equal(t, "A3", result.state.Messages[3].Content) + + rollbackEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventRollback + }) + require.Len(t, rollbackEvents, 1) + require.NotNil(t, rollbackEvents[0].Rollback) + assert.Equal(t, t1.EventID, rollbackEvents[0].Rollback.ToEventID) + assert.Equal(t, t2.EventID, rollbackEvents[0].Rollback.PreviousHeadCommitEventID) + assert.NotContains(t, store.checkpoints, sessionRunnerCheckpointID(sid)) +} + +func TestRollbackSessionMultipleRollbacksProjectActiveBranch(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "rollback-multiple" + + appendCommittedTestTurn(t, ctx, store, sid, "turn-1", "Q1", "A1") + appendCommittedTestTurn(t, ctx, store, sid, "turn-2", "Q2", "A2") + require.NoError(t, RollbackSession[*schema.Message](ctx, store, sid, "turn-1")) + appendCommittedTestTurn(t, ctx, store, sid, "turn-3", "Q3", "A3") + require.NoError(t, RollbackSession[*schema.Message](ctx, store, sid, "turn-1")) + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 2) + assert.Equal(t, "Q1", result.state.Messages[0].Content) + assert.Equal(t, "A1", result.state.Messages[1].Content) + + rollbackEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventRollback + }) + require.Len(t, rollbackEvents, 2) +} + +func TestRunnerQueryAfterRollbackUsesActiveProjection(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "runner-query-after-rollback" + + firstAgent := &runnerSessionAgent{ + name: "runner-session-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("first"), schema.AssistantMessage("answer1", nil)}, + }, + } + firstRunner := NewRunner(ctx, RunnerConfig{ + Agent: firstAgent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, firstRunner.Query(ctx, "first")) + firstCommittedIdleEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return isCommittedIdleEvent(se) + }) + require.Len(t, firstCommittedIdleEvents, 1) + firstBoundaryEventID := firstCommittedIdleEvents[0].EventID + + secondAgent := &runnerSessionAgent{ + name: "runner-session-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("first"), schema.AssistantMessage("answer1", nil), schema.UserMessage("second"), schema.AssistantMessage("answer2", nil)}, + }, + } + secondRunner := NewRunner(ctx, RunnerConfig{ + Agent: secondAgent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, secondRunner.Query(ctx, "second")) + + require.NoError(t, RollbackSession[*schema.Message](ctx, store, sid, firstBoundaryEventID)) + + thirdAgent := &runnerSessionAgent{ + name: "runner-session-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("first"), schema.AssistantMessage("answer1", nil), schema.UserMessage("third"), schema.AssistantMessage("answer3", nil)}, + }, + } + thirdRunner := NewRunner(ctx, RunnerConfig{ + Agent: thirdAgent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, thirdRunner.Query(ctx, "third")) + + require.Len(t, thirdAgent.inputs, 1) + require.Len(t, thirdAgent.inputs[0], 3) + assert.Equal(t, "first", thirdAgent.inputs[0][0].Content) + assert.Equal(t, "ok", thirdAgent.inputs[0][1].Content) + assert.Equal(t, "third", thirdAgent.inputs[0][2].Content) +} + +func TestRollbackSessionTargetResolutionErrors(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "rollback-target-errors" + + appendCommittedTestTurn(t, ctx, store, sid, "turn-1", "Q1", "A1") + appendCommittedTestTurn(t, ctx, store, sid, "turn-2", "Q2", "A2") + appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + EventID: "pending-event", + Kind: SessionEventMessage, + Message: testMessageWithID("pending", schema.User), + }) + + err := RollbackSession[*schema.Message](ctx, store, sid, "pending-event") + require.ErrorIs(t, err, ErrInvalidRollbackTarget) + + err = RollbackSession[*schema.Message](ctx, store, sid, "missing") + require.ErrorIs(t, err, ErrRollbackTargetNotFound) + + err = RollbackSession[*schema.Message]( + ctx, + store, + sid, + "turn-1", + WithRollbackSessionExpectedHeadEventID[*schema.Message]("stale-head"), + ) + require.ErrorIs(t, err, ErrSessionHeadChanged) + rollbackEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventRollback + }) + require.Empty(t, rollbackEvents) + + require.NoError(t, RollbackSession[*schema.Message]( + ctx, + store, + sid, + "turn-1", + WithRollbackSessionExpectedHeadEventID[*schema.Message]("turn-2"), + )) + err = RollbackSession[*schema.Message]( + ctx, + store, + sid, + "turn-2", + WithRollbackSessionExpectedHeadEventID[*schema.Message]("turn-2"), + ) + require.ErrorIs(t, err, ErrRollbackTargetInactive) +} + +func TestReconstructRollbackMalformedRecordsFailClosed(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "rollback-malformed" + + msg := appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Message: testMessageWithID("Q1", schema.User), + }) + appendCommittedTestTurn(t, ctx, store, sid, "turn-1", "A1") + + appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + Kind: SessionEventRollback, + Rollback: &SessionRollbackEvent{ + ToEventID: msg.EventID, + }, + }) + _, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.ErrorIs(t, err, ErrInvalidRollbackTarget) + + store = newSessionHelperStore() + appendCommittedTestTurn(t, ctx, store, sid, "turn-1", "Q1", "A1") + payloadEvent := &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Kind: SessionEventRollback, + Rollback: &SessionRollbackEvent{ + ToEventID: "missing-turn-end-event", + }, + } + data, encodeErr := encodeSessionEvent(payloadEvent) + require.NoError(t, encodeErr) + store.mu.Lock() + store.events = append(store.events, storedSessionEvent{ + EventID: payloadEvent.EventID, + Kind: payloadEvent.Kind, + Data: append([]byte{}, data...), + }) + store.eventIDs = append(store.eventIDs, payloadEvent.EventID) + store.eventIDIdx[payloadEvent.EventID] = len(store.events) - 1 + store.mu.Unlock() + _, err = reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.ErrorIs(t, err, ErrRollbackTargetInactive) + + store = newSessionHelperStore() + appendCommittedTestTurn(t, ctx, store, sid, "turn-1", "Q1", "A1") + staleTarget := appendCommittedTestTurn(t, ctx, store, sid, "turn-2", "Q2", "A2") + require.NoError(t, RollbackSession[*schema.Message](ctx, store, sid, "turn-1")) + appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + Kind: SessionEventRollback, + Rollback: &SessionRollbackEvent{ + ToEventID: staleTarget.EventID, + }, + }) + _, err = reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.ErrorIs(t, err, ErrRollbackTargetInactive) +} + +// TestRunnerSessionReconstructsFromEventLog: Delete testTurnState from store, +// next turn should reconstruct from events. +func TestRunnerSessionReconstructsFromEventLog(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "reconstruct-session" + + firstAgent := &runnerSessionAgent{ + name: "ra", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("first"), schema.AssistantMessage("answer1", nil)}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: firstAgent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "first")) + + // Verify context-commit events were captured: caller input + assistant output + turn-end. + commitEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage || isCommittedIdleEvent(se) + }) + require.Len(t, commitEvents, 3, "input event + assistant event + turn-end event should be in event log") + + // Capture the prepared session state before agent runs. + capturedAgent := &runnerSessionAgent{ + name: "ra", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{}, + }, + } + runner = NewRunner(ctx, RunnerConfig{ + Agent: capturedAgent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "second")) + + // The agent should have received the reconstructed history before "second". + require.Len(t, capturedAgent.inputs, 1) + // Input order: reconstructed user "first" + reconstructed assistant "ok" + new user "second". + require.Len(t, capturedAgent.inputs[0], 3) + // The last message must be the new "second" input. + assert.Equal(t, "second", capturedAgent.inputs[0][len(capturedAgent.inputs[0])-1].Content) + // And the first reconstructed message must be the original "first" input. + assert.Equal(t, "first", capturedAgent.inputs[0][0].Content) +} + +// TestRunnerSessionInputEventsPersisted verifies that caller input messages +// are persisted to the event log at turn start. +func TestRunnerSessionInputEventsPersisted(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "input-events" + + agent := &runnerSessionAgent{ + name: "input-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("answer", nil)}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "user-question")) + + // Single-turn run: 1 user input event + 1 assistant output event + 1 idle commit event, + // plus non-context lifecycle timeline records. + commitEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage || isCommittedIdleEvent(se) + }) + require.Len(t, commitEvents, 3) + // The first message event should be the user input. + first := commitEvents[0] + require.NotNil(t, first.Message) + assert.Equal(t, "user-question", first.Message.Content) + assert.Equal(t, schema.User, first.Message.Role) + // And it should have a message ID. + assert.NotEmpty(t, GetMessageID(first.Message)) +} + +// recordingHelperStore wraps sessionHelperStore to record the order of +// AppendEvents and Set calls so tests can assert durability ordering. +type recordingHelperStore struct { + *sessionHelperStore + mu sync.Mutex + calls []string // "append" or "set:" + delaySet time.Duration +} + +func newRecordingHelperStore() *recordingHelperStore { + return &recordingHelperStore{sessionHelperStore: newSessionHelperStore()} +} + +func (s *recordingHelperStore) AppendEventsForSession(ctx context.Context, sid string, events []*SessionEvent[*schema.Message]) error { + s.mu.Lock() + if s.sessionHelperStore.appendErr != nil { + err := s.sessionHelperStore.appendErr + s.mu.Unlock() + return err + } + s.calls = append(s.calls, "append") + s.mu.Unlock() + return s.sessionHelperStore.AppendEventsForSession(ctx, sid, events) +} + +func (s *recordingHelperStore) AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, sessionID, events) +} + +func (s *recordingHelperStore) openSession(_ context.Context, req *openSessionRequest) (*openSessionResult[*schema.Message], error) { + sessionID := "" + if req != nil { + sessionID = req.sessionID + } + return &openSessionResult[*schema.Message]{handle: &legacyMessageTestHandle{store: s, sessionID: sessionID}}, nil +} + +func (s *recordingHelperStore) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, "", events) +} + +func (s *recordingHelperStore) Set(ctx context.Context, key string, value []byte) error { + if s.delaySet > 0 { + time.Sleep(s.delaySet) + } + s.mu.Lock() + s.calls = append(s.calls, "set:"+key) + s.mu.Unlock() + return s.sessionHelperStore.Set(ctx, key, value) +} + +func (s *recordingHelperStore) callsSnapshot() []string { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]string, len(s.calls)) + copy(out, s.calls) + return out +} + +// TestRunnerSessionInterruptCheckpointSkippedOnPersistFailure proves the +// fail-closed invariant: if AppendEvents fails during a turn that ends in an +// interrupt, the checkpoint MUST NOT be written — otherwise resume would load +// a checkpoint referencing events that were never persisted. +func TestRunnerSessionInterruptCheckpointSkippedOnPersistFailure(t *testing.T) { + ctx := context.Background() + store := newRecordingHelperStore() + store.sessionHelperStore.kindErr = map[SessionEventKind]error{ + SessionEventInterrupt: errors.New("simulated append failure"), + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerInterruptAgent{}, + CheckPointStore: store, + SessionID: "interrupt-persist-fail", + SessionStore: store, + }) + iter := runner.Query(ctx, "go") + var sawErr bool + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + sawErr = true + } + } + require.True(t, sawErr, "expected runner to surface the persistence error") + + cpKey := sessionRunnerCheckpointID("interrupt-persist-fail") + calls := store.callsSnapshot() + for _, c := range calls { + if c == "set:"+cpKey { + t.Fatalf("checkpoint was written despite event persistence failure: calls=%v", calls) + } + } +} + +// TestRunnerSessionCheckpointAfterPersisterFlush proves that on the interrupt +// path, the checkpoint is written ONLY after the persister has flushed events +// (AppendEvents before Set on checkpoint key). +func TestRunnerSessionCheckpointAfterPersisterFlush(t *testing.T) { + ctx := context.Background() + store := newRecordingHelperStore() + + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerInterruptAgent{}, + CheckPointStore: store, + SessionID: "interrupt-order", + SessionStore: store, + }) + iter := runner.Query(ctx, "hi") + for { + _, ok := iter.Next() + if !ok { + break + } + } + calls := store.callsSnapshot() + + cpKey := sessionRunnerCheckpointID("interrupt-order") + var lastAppend, firstSet int = -1, -1 + for i, c := range calls { + if c == "append" { + lastAppend = i + } + if c == "set:"+cpKey && firstSet == -1 { + firstSet = i + } + } + require.NotEqual(t, -1, lastAppend, "expected at least one AppendEvents call") + require.NotEqual(t, -1, firstSet, "expected the runner-session checkpoint to be written") + require.Greater(t, firstSet, lastAppend, + "checkpoint Set must follow the final AppendEvents flush; got calls=%v", calls) +} + +func TestRunnerSessionInterruptCheckpointTailIsFinalIdle(t *testing.T) { + ctx := context.Background() + store := newRecordingHelperStore() + sid := "interrupt-tail" + + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerInterruptAgent{}, + CheckPointStore: store, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "hi")) + + cpKey := sessionRunnerCheckpointID(sid) + raw, ok := store.checkpoints[cpKey] + require.True(t, ok, "expected interrupt checkpoint to be saved") + cp, err := decodeRunnerSessionCheckpoint(raw) + require.NoError(t, err) + + store.sessionHelperStore.mu.Lock() + require.NotEmpty(t, store.events) + tail := store.events[len(store.events)-1] + store.sessionHelperStore.mu.Unlock() + assert.Equal(t, SessionEventSessionStatusIdle, tail.Kind) + + _, runCtx, _, err := runnerLoadCheckPointBytes(ctx, cp.Payload) + require.NoError(t, err) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + for _, event := range runCtx.Session.Events { + require.NotNil(t, event.AgentEvent) + assert.Nil(t, event.SessionEventVariant) + } +} + +func TestRunnerSessionCheckpointPayloadStripsSessionEvents(t *testing.T) { + ctx := context.Background() + store := newRecordingHelperStore() + sid := "checkpoint-strip-session-events" + + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerCheckpointSanitizeAgent{}, + CheckPointStore: store, + SessionID: sid, + SessionStore: store, + }) + iter := runner.Query(ctx, "hi", WithTimelineEvents()) + var liveSessionEventIDs []string + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil { + liveSessionEventIDs = append(liveSessionEventIDs, event.SessionEventVariant.Event.EventID) + } + } + assert.Contains(t, liveSessionEventIDs, "checkpoint-session-only") + assert.Contains(t, liveSessionEventIDs, "checkpoint-output") + + cpKey := sessionRunnerCheckpointID(sid) + raw, ok := store.checkpoints[cpKey] + require.True(t, ok, "expected interrupt checkpoint to be saved") + cp, err := decodeRunnerSessionCheckpoint(raw) + require.NoError(t, err) + + _, runCtx, _, err := runnerLoadCheckPointBytes(ctx, cp.Payload) + require.NoError(t, err) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + + var foundOutput bool + for _, event := range runCtx.Session.Events { + require.NotNil(t, event.AgentEvent) + assert.Nil(t, event.SessionEventVariant) + assert.True(t, event.Output != nil || event.Action != nil || event.Err != nil) + if event.Output != nil && + event.Output.MessageOutput != nil && + event.Output.MessageOutput.Message != nil && + event.Output.MessageOutput.Message.Content == "mixed output" { + foundOutput = true + } + } + assert.True(t, foundOutput) + + var persistedKinds []SessionEventKind + store.sessionHelperStore.mu.Lock() + for _, event := range store.events { + persistedKinds = append(persistedKinds, event.Kind) + } + store.sessionHelperStore.mu.Unlock() + assert.Contains(t, persistedKinds, SessionEventSessionStatusRunning) + assert.Contains(t, persistedKinds, SessionEventMessage) +} + +func TestRunnerSessionAgentInterruptBoundaryFailureNotExposed(t *testing.T) { + ctx := context.Background() + store := newRecordingHelperStore() + store.sessionHelperStore.kindErr = map[SessionEventKind]error{ + SessionEventInterrupt: errors.New("agent interrupt append failed"), + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerInterruptAgent{}, + CheckPointStore: store, + SessionID: "interrupt-not-exposed", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hi", WithTimelineEvents()) + var kinds []SessionEventKind + var errs []error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + errs = append(errs, event.Err) + } + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil { + kinds = append(kinds, event.SessionEventVariant.Event.Kind) + } + } + require.NotEmpty(t, errs) + assert.NotContains(t, kinds, SessionEventInterrupt) + + cpKey := sessionRunnerCheckpointID("interrupt-not-exposed") + _, existed := store.checkpoints[cpKey] + assert.False(t, existed, "checkpoint must not be saved after interrupt boundary append failure") +} + +func TestRunnerSessionInterruptPersistErrorSurfacesWithoutCheckpoint(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + store.kindErr = map[SessionEventKind]error{ + SessionEventInterrupt: errors.New("agent interrupt append failed"), + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: &runnerInterruptAgent{}, + SessionID: "interrupt-no-checkpoint", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hi") + var errs []error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + errs = append(errs, event.Err) + } + } + require.NotEmpty(t, errs) + assert.ErrorContains(t, errs[len(errs)-1], "failed to persist session events") +} + +// TestSessionPersister_EnqueueAfterAppendError verifies that once AppendEvents +// has failed, subsequent enqueue calls return that error rather than silently +// succeeding. +func TestSessionPersister_EnqueueAfterAppendError(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + store.appendErr = errors.New("append failed") + + p := newSessionEventPersister[*schema.Message](ctx, store, "sid") + + require.NoError(t, p.enqueueAsync(validTestPayload())) + require.Error(t, p.closeAndWait()) + require.Error(t, p.getErr(), "persister must record the AppendEvents failure") + + err := p.enqueueAsync(validTestPayload()) + require.Error(t, err, "enqueue after persist failure must return an error") + assert.Contains(t, err.Error(), "append failed") + + for i := 0; i < 4; i++ { + err = p.enqueueAsync(validTestPayload()) + require.Error(t, err, "latched error must be returned consistently") + assert.Contains(t, err.Error(), "append failed") + } +} + +// transientFailStore fails the first N AppendEvents calls then succeeds. +type transientFailStore struct { + sessionHelperStore + retryMu sync.Mutex + failsLeft int + appendCalls int + appendErrVal error +} + +func (s *transientFailStore) AppendEventsForSession(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + s.retryMu.Lock() + s.appendCalls++ + if s.failsLeft > 0 { + s.failsLeft-- + s.retryMu.Unlock() + return s.appendErrVal + } + s.retryMu.Unlock() + return s.sessionHelperStore.AppendEventsForSession(ctx, sessionID, events) +} + +func (s *transientFailStore) AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, sessionID, events) +} + +func (s *transientFailStore) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return s.AppendEventsForSession(ctx, "", events) +} + +func (s *transientFailStore) getAppendCalls() int { + s.retryMu.Lock() + defer s.retryMu.Unlock() + return s.appendCalls +} + +func TestSessionPersister_FlushDoesNotRetryTransientFailure(t *testing.T) { + ctx := context.Background() + store := &transientFailStore{ + sessionHelperStore: *newSessionHelperStore(), + failsLeft: 2, + appendErrVal: errors.New("transient"), + } + + p := newSessionEventPersister[*schema.Message](ctx, store, "sid") + + require.NoError(t, p.enqueueAsync(validTestPayload())) + + err := p.closeAndWait() + require.Error(t, err) + assert.Contains(t, err.Error(), "transient") + assert.Equal(t, 1, store.getAppendCalls()) + store.sessionHelperStore.mu.Lock() + assert.Empty(t, store.sessionHelperStore.events) + store.sessionHelperStore.mu.Unlock() +} + +func TestSessionPersister_FlushPermanentFailureLatched(t *testing.T) { + ctx := context.Background() + store := &transientFailStore{ + sessionHelperStore: *newSessionHelperStore(), + failsLeft: 100, // always fail + appendErrVal: errors.New("permanent"), + } + + p := newSessionEventPersister[*schema.Message](ctx, store, "sid") + + require.NoError(t, p.enqueueAsync(validTestPayload())) + + err := p.closeAndWait() + require.Error(t, err) + assert.Contains(t, err.Error(), "permanent") + assert.Equal(t, 1, store.getAppendCalls()) +} + +func TestSessionPersister_FlushContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + store := &transientFailStore{ + sessionHelperStore: *newSessionHelperStore(), + failsLeft: 100, // always fail + appendErrVal: errors.New("failing"), + } + + p := newSessionEventPersister[*schema.Message](ctx, store, "sid") + + require.NoError(t, p.enqueueAsync(validTestPayload())) + + cancel() + + err := p.closeAndWait() + require.Error(t, err) + assert.Contains(t, err.Error(), "failing") + assert.Equal(t, 1, store.getAppendCalls()) +} + +// --- Attack tests for reconstruction across committed and in-flight events --- + +// TestAttack_ReconstructionIncludesInterruptedTailOnResume verifies that +// reconstructSessionState keeps interrupted-tail messages during replay. +func TestAttack_ReconstructionIncludesInterruptedTailOnResume(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "inflight-recovery" + + committedMsg := schema.UserMessage("committed-msg") + EnsureMessageID(committedMsg) + + // A committed turn: TurnStart (lifecycle running) + Message + committed idle. + events := []*SessionEvent[*schema.Message]{ + {EventID: uuid.NewString(), Kind: SessionEventSessionStatusRunning, Lifecycle: &LifecycleEvent{State: SessionRunStateRunning}}, + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: committedMsg}, + {EventID: uuid.NewString(), Kind: SessionEventSessionStatusIdle, Lifecycle: &LifecycleEvent{State: SessionRunStateIdle, StopReason: &StopReason{Type: "end_turn"}}}, + } + + // An interrupted run: a Message event with no committed idle. + interruptedMsg := schema.AssistantMessage("interrupted-msg", nil) + EnsureMessageID(interruptedMsg) + events = append(events, &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), Kind: SessionEventMessage, Message: interruptedMsg, + }) + + for _, se := range events { + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + // State should have messages from committed turn (1) + interrupted turn (1). + require.Len(t, result.state.Messages, 2) + assert.Equal(t, "committed-msg", result.state.Messages[0].Content) + assert.Equal(t, "interrupted-msg", result.state.Messages[1].Content) +} + +func TestAttack_ReconstructionWithoutCommittedIdle(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "inflight-no-committed-turn" + + msg := schema.UserMessage("first-turn") + EnsureMessageID(msg) + events := []*SessionEvent[*schema.Message]{ + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: msg}, + {EventID: uuid.NewString(), Kind: SessionEventInterrupt, Interrupt: &InterruptEvent{ + Contexts: []*InterruptContext{ + { + InterruptID: "agent:InterruptAgent", + Info: "approval_needed", + }, + }, + }}, + } + for _, se := range events { + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 1) + assert.Equal(t, "first-turn", result.state.Messages[0].Content) +} + +// TestAttack_OldRunIDFieldIgnoredOnDeserialization verifies that a JSON payload +// containing a legacy "run_id" field is deserialized without error. +func TestAttack_OldRunIDFieldIgnoredOnDeserialization(t *testing.T) { + rawJSON := []byte(`{ + "event_id": "evt-legacy", + "run_id": "old-run", + "kind": "message", + "message": {"role": "user", "content": "hello from legacy"} + }`) + + event, err := decodeSessionEventWithSerializer[*schema.Message](rawJSON, nil) + require.NoError(t, err, "deserialization must not fail on unknown run_id field") + require.NotNil(t, event) + assert.Equal(t, "evt-legacy", event.EventID) + require.NotNil(t, event.Message) + assert.Equal(t, "hello from legacy", event.Message.Content) +} + +func TestAttack_ResumeAfterInterruptedRunWritesSessionEvents(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sessionID := "resume-after-interrupt" + + // First, run a normal turn that completes (provides a committed idle baseline). + normalAgent := &runnerSessionAgent{ + name: "normal-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("first answer", nil)}, + }, + } + firstRunner := NewRunner(ctx, RunnerConfig{ + Agent: normalAgent, + SessionID: sessionID, + SessionStore: store, + CheckPointStore: store, + }) + drainSessionEvents(t, firstRunner.Query(ctx, "first question")) + + // Now run a query that interrupts (building on the committed session). + agent := &runnerInterruptAgent{} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sessionID, + SessionStore: store, + CheckPointStore: store, + }) + + iter := runner.Query(ctx, "trigger interrupt") + for { + _, ok := iter.Next() + if !ok { + break + } + } + + eventsBeforeResume := len(store.events) + + resumeIter, err := runner.Resume(ctx, "") + require.NoError(t, err) + for { + _, ok := resumeIter.Next() + if !ok { + break + } + } + + require.Greater(t, len(store.events), eventsBeforeResume) + resumeEvents := filterStoredSessionEvents(t, store.events[eventsBeforeResume:], func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventKind(SessionEventExtensionPrefix+"resume.request_started") || + se.Kind == SessionEventSessionStatusIdle + }) + require.NotEmpty(t, resumeEvents) +} + +func TestAttack_FreshRunIgnoresInterruptedSuffixMetadata(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sessionID := "fresh-run-ignores-inflight" + + // First, run a normal turn that completes (provides a committed idle baseline). + normalAgent := &runnerSessionAgent{ + name: "normal-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("baseline", nil)}, + }, + } + baselineRunner := NewRunner(ctx, RunnerConfig{ + Agent: normalAgent, + SessionID: sessionID, + SessionStore: store, + CheckPointStore: store, + }) + drainSessionEvents(t, baselineRunner.Query(ctx, "baseline")) + + // Now run a query that interrupts. + agent := &runnerInterruptAgent{} + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sessionID, + SessionStore: store, + CheckPointStore: store, + }) + + iter := runner.Query(ctx, "trigger interrupt") + for { + _, ok := iter.Next() + if !ok { + break + } + } + + // Instead of resuming, create a NEW runner on the same session and run a new query (fresh Run). + eventsBeforeFresh := len(store.events) + freshAgent := &runnerSessionAgent{ + name: "fresh-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("fresh answer", nil)}, + }, + } + freshRunner := NewRunner(ctx, RunnerConfig{ + Agent: freshAgent, + SessionID: sessionID, + SessionStore: store, + CheckPointStore: store, + }) + drainSessionEvents(t, freshRunner.Query(ctx, "new question")) + + require.Greater(t, len(store.events), eventsBeforeFresh) + require.Len(t, freshAgent.inputs, 1) + require.Len(t, freshAgent.inputs[0], 4) + assert.Equal(t, "baseline", freshAgent.inputs[0][0].Content) + assert.Equal(t, "ok", freshAgent.inputs[0][1].Content) + assert.Equal(t, "trigger interrupt", freshAgent.inputs[0][2].Content) + assert.Equal(t, "new question", freshAgent.inputs[0][3].Content) +} + +// sessionStreamingAgent emits a single streaming assistant output. Used to +// verify the runner's stream-copy/persist path. +type sessionStreamingAgent struct { + chunks []*schema.Message + streamErr error + turnEnd *testTurnState[*schema.Message] + role schema.RoleType + tool string + preEvent *SessionEvent[*schema.Message] +} + +func (a *sessionStreamingAgent) Name(_ context.Context) string { return "session-stream-agent" } +func (a *sessionStreamingAgent) Description(_ context.Context) string { return "stream test agent" } +func (a *sessionStreamingAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + if a.preEvent != nil { + gen.Send(&AgentEvent{ + AgentName: "session-stream-agent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: a.preEvent, + }, + }) + } + stream := testStreamReaderWithTerminalError(a.chunks, a.streamErr) + role := a.role + if role == "" { + role = schema.Assistant + } + mv := &MessageVariant{IsStreaming: true, MessageStream: stream, Role: role, ToolName: a.tool} + gen.Send(&AgentEvent{AgentName: "session-stream-agent", Output: &AgentOutput{MessageOutput: mv}}) + }() + return iter +} + +type agenticSessionStreamingAgent struct { + chunks []*schema.AgenticMessage + streamErr error + turnEnd *testTurnState[*schema.AgenticMessage] +} + +func (a *agenticSessionStreamingAgent) Name(_ context.Context) string { + return "agentic-session-stream-agent" +} + +func (a *agenticSessionStreamingAgent) Description(_ context.Context) string { + return "agentic stream test agent" +} + +func (a *agenticSessionStreamingAgent) Run( + _ context.Context, + _ *TypedAgentInput[*schema.AgenticMessage], + _ ...AgentRunOption, +) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer gen.Close() + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "agentic-session-stream-agent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: testStreamReaderWithTerminalError(a.chunks, a.streamErr), + AgenticRole: schema.AgenticRoleTypeUser, + }, + }, + }) + }() + return iter +} + +func testStreamReaderWithTerminalError[T any](chunks []T, streamErr error) *schema.StreamReader[T] { + if streamErr == nil { + return schema.StreamReaderFromArray(chunks) + } + reader, writer := schema.Pipe[T](len(chunks) + 1) + go func() { + defer writer.Close() + for _, chunk := range chunks { + writer.Send(chunk, nil) + } + var zero T + writer.Send(zero, streamErr) + }() + return reader +} + +// TestStreamPersistence_CopyAndConcat verifies that streaming assistant outputs +// produce a durable, fully-concatenated SessionEvent.Message AND remain consumable +// from the live stream. Regression test for the pre-evaluation bug where +// stream-only events (Message==nil, MessageStream!=nil) skipped persistence. +func TestStreamPersistence_CopyAndConcat(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "stream-session" + + chunks := []*schema.Message{ + schema.AssistantMessage("hello ", nil), + schema.AssistantMessage("world", nil), + } + agent := &sessionStreamingAgent{ + chunks: chunks, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("q"), schema.AssistantMessage("hello world", nil)}, + }, + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + // Drain live events and verify the live stream still produces the concatenated content. + iter := runner.Query(ctx, "q") + var liveContent string + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + msg, err := schema.ConcatMessageStream(ev.Output.MessageOutput.MessageStream) + require.NoError(t, err) + liveContent = msg.Content + } + } + assert.Equal(t, "hello world", liveContent, "live stream must yield concatenated content") + + // Find the persisted streaming event in the log: exactly one assistant output should be persisted. + var assistantMessages []*schema.Message + for _, ep := range store.events { + se, err := decodeSessionEvent[*schema.Message](ep.Data) + require.NoError(t, err) + if se.Message != nil && se.Message.Role == schema.Assistant { + assistantMessages = append(assistantMessages, se.Message) + } + } + require.Len(t, assistantMessages, 1, "streaming assistant output must be persisted exactly once") + assert.Equal(t, "hello world", assistantMessages[0].Content, + "persisted stream message must be the fully concatenated content") +} + +func TestStreamPersistence_IncompleteStreamPrefixPersisted(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + streamErr := errors.New("model stream failed") + agent := &sessionStreamingAgent{ + chunks: []*schema.Message{ + schema.AssistantMessage("hello ", nil), + schema.AssistantMessage("partial", nil), + }, + streamErr: streamErr, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: "incomplete-stream-session", + SessionStore: store, + }) + + drainErroredStreamEvents(t, runner.Query(ctx, "q"), streamErr) + + events := decodeStoredSessionEvents(t, store.events) + var incomplete []*SessionEvent[*schema.Message] + var normalFailedMessages []*SessionEvent[*schema.Message] + for _, se := range events { + if se.Kind == SessionEventMessageStreamIncomplete { + incomplete = append(incomplete, se) + } + if se.Kind == SessionEventMessage && se.Message != nil && + se.Message.Role == schema.Assistant && se.Message.Content == "hello partial" { + normalFailedMessages = append(normalFailedMessages, se) + } + } + require.Len(t, incomplete, 1) + require.NotNil(t, incomplete[0].MessageStreamIncomplete) + require.NotNil(t, incomplete[0].MessageStreamIncomplete.Message) + assert.Equal(t, "hello partial", incomplete[0].MessageStreamIncomplete.Message.Content) + assert.Contains(t, incomplete[0].MessageStreamIncomplete.Error, streamErr.Error()) + assert.Empty(t, normalFailedMessages, "failed stream prefix must not be persisted as a normal context message") +} + +func TestAttack_IncompleteStreamPrefixCarriesDurableMetadata(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "attack-incomplete-metadata" + streamErr := errors.New("stream transport failed") + agent := &sessionStreamingAgent{ + chunks: []*schema.Message{ + schema.AssistantMessage("prefix", nil), + }, + streamErr: streamErr, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + drainErroredStreamEvents(t, runner.Query(ctx, "q"), streamErr) + + events := decodeStoredSessionEvents(t, store.events) + var incomplete *SessionEvent[*schema.Message] + var idle *SessionEvent[*schema.Message] + for _, se := range events { + switch se.Kind { + case SessionEventMessageStreamIncomplete: + incomplete = se + case SessionEventSessionStatusIdle: + idle = se + } + } + + require.NotNil(t, incomplete) + require.NotNil(t, idle) + assert.NotEmpty(t, incomplete.EventID) + assert.True(t, incomplete.Timestamp.Before(idle.Timestamp) || incomplete.Timestamp.Equal(idle.Timestamp)) + assert.Equal(t, "prefix", incomplete.MessageStreamIncomplete.Message.Content) + assert.Contains(t, incomplete.MessageStreamIncomplete.Error, streamErr.Error()) +} + +func TestAttack_IncompleteStreamPersistsAllTerminalErrors(t *testing.T) { + ctx := context.Background() + tests := []struct { + name string + streamErr error + }{ + {name: "canceled", streamErr: ErrStreamCanceled}, + {name: "will retry", streamErr: &WillRetryError{ErrStr: "retry", RetryAttempt: 1}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: &sessionStreamingAgent{ + chunks: []*schema.Message{schema.AssistantMessage("transient", nil)}, + streamErr: tt.streamErr, + }, + EnableStreaming: true, + SessionID: "attack-nondurable-" + tt.name, + SessionStore: store, + }) + + drainErroredStreamEvents(t, runner.Query(ctx, "q"), tt.streamErr) + + incomplete := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessageStreamIncomplete + }) + require.Len(t, incomplete, 1) + require.NotNil(t, incomplete[0].MessageStreamIncomplete) + assert.Equal(t, "transient", incomplete[0].MessageStreamIncomplete.Message.Content) + assert.Contains(t, incomplete[0].MessageStreamIncomplete.Error, tt.streamErr.Error()) + }) + } +} + +func drainErroredStreamEvents(t *testing.T, iter *AsyncIterator[*AgentEvent], streamErr error) { + t.Helper() + var sawStreamErr bool + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output == nil || ev.Output.MessageOutput == nil || + !ev.Output.MessageOutput.IsStreaming || ev.Output.MessageOutput.MessageStream == nil { + continue + } + for { + _, err := ev.Output.MessageOutput.MessageStream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + assert.ErrorContains(t, err, streamErr.Error()) + sawStreamErr = true + break + } + } + } + require.True(t, sawStreamErr, "live stream must surface the terminal stream error") +} + +func TestStreamPersistence_IncompleteStreamExcludedFromReconstruction(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "incomplete-reconstruct-session" + appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Message: schema.UserMessage("q"), + }) + appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + Kind: SessionEventMessageStreamIncomplete, + MessageStreamIncomplete: &MessageStreamIncompleteEvent[*schema.Message]{ + Message: schema.AssistantMessage("partial", nil), + Error: "model stream failed", + }, + }) + appendTestSessionEvent(t, ctx, store, sid, &SessionEvent[*schema.Message]{ + Kind: SessionEventSessionStatusIdle, + Lifecycle: &LifecycleEvent{ + State: SessionRunStateIdle, + StopReason: &StopReason{Type: "end_turn"}, + }, + }) + + result, err := reconstructSessionState[*schema.Message](ctx, mustOpenTestSession[*schema.Message](t, ctx, store, sid), sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 1) + assert.Equal(t, "q", result.state.Messages[0].Content) +} + +func TestMessageStreamIncompleteEvent_RoundTripAndValidation(t *testing.T) { + event := withTestEventID(&SessionEvent[*schema.Message]{ + Kind: SessionEventMessageStreamIncomplete, + MessageStreamIncomplete: &MessageStreamIncompleteEvent[*schema.Message]{ + Message: schema.AssistantMessage("partial", nil), + Error: "model stream failed", + }, + }) + encoded, err := encodeSessionEvent(event) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](encoded) + require.NoError(t, err) + require.NotNil(t, decoded.MessageStreamIncomplete) + assert.Equal(t, SessionEventMessageStreamIncomplete, decoded.Kind) + assert.Equal(t, "partial", decoded.MessageStreamIncomplete.Message.Content) + assert.Equal(t, "model stream failed", decoded.MessageStreamIncomplete.Error) + assert.False(t, isContextSessionEvent(decoded)) + + _, err = encodeSessionEvent(withTestEventID(&SessionEvent[*schema.Message]{ + Kind: SessionEventMessageStreamIncomplete, + MessageStreamIncomplete: &MessageStreamIncompleteEvent[*schema.Message]{Error: "missing message"}, + })) + require.Error(t, err) + assert.Contains(t, err.Error(), "message stream incomplete event") +} + +func TestStreamPersistence_StreamingLiveBeforeMaterializedBoundary(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "sync-stream-session" + + agent := &sessionStreamingAgent{ + chunks: []*schema.Message{ + schema.AssistantMessage("hello ", nil), + schema.AssistantMessage("sync", nil), + }, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("q"), schema.AssistantMessage("hello sync", nil)}, + }, + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + iter := runner.Query(ctx, "q") + var observed *MessageVariant + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil { + observed = ev.Output.MessageOutput + } + } + + require.NotNil(t, observed) + assert.True(t, observed.IsStreaming, "streaming output remains live while persistence materializes a copy") + msg, err := observed.GetMessage() + require.NoError(t, err) + assert.Equal(t, "hello sync", msg.Content) + + var stored bool + store.mu.Lock() + snapshot := append([]storedSessionEvent{}, store.events...) + store.mu.Unlock() + for _, ep := range snapshot { + se, err := decodeSessionEvent[*schema.Message](ep.Data) + require.NoError(t, err) + if se.Message != nil && se.Message.Role == schema.Assistant && se.Message.Content == "hello sync" { + stored = true + } + } + assert.True(t, stored, "materialized stream message must be persisted by finalization") +} + +func TestStreamPersistence_PendingAnnotationFlushesBeforeMaterializedBoundary(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + annotationKind := SessionEventKind(SessionEventExtensionPrefix + "stream.annotation") + agent := &sessionStreamingAgent{ + preEvent: &SessionEvent[*schema.Message]{ + Kind: annotationKind, + Extension: &SessionExtensionEvent{}, + }, + chunks: []*schema.Message{ + schema.AssistantMessage("hello ", nil), + schema.AssistantMessage("stream", nil), + }, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("q"), schema.AssistantMessage("hello stream", nil)}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: "stream-annotation-boundary", + SessionStore: store, + }) + + drainSessionEvents(t, runner.Query(ctx, "q")) + + assert.Equal(t, [][]SessionEventKind{ + {SessionEventSessionStatusRunning}, + {SessionEventMessage}, + {annotationKind}, + {SessionEventMessage}, + {SessionEventSessionStatusIdle}, + }, store.appendBatches) +} + +func TestStreamPersistence_ToolResultStreamingLiveBeforeMaterializedBoundary(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "sync-tool-stream-session" + + agent := &sessionStreamingAgent{ + chunks: []*schema.Message{ + schema.ToolMessage("tool ", "tc-1", schema.WithToolName("t1")), + schema.ToolMessage("result", "tc-1", schema.WithToolName("t1")), + }, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.ToolMessage("tool result", "tc-1", schema.WithToolName("t1"))}, + }, + role: schema.Tool, + tool: "t1", + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + iter := runner.Query(ctx, "q") + var observed *MessageVariant + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil { + observed = ev.Output.MessageOutput + } + } + + require.NotNil(t, observed) + assert.True(t, observed.IsStreaming) + msg, err := observed.GetMessage() + require.NoError(t, err) + assert.Equal(t, schema.Tool, msg.Role) + assert.Equal(t, "tool result", msg.Content) + + var stored bool + store.mu.Lock() + snapshot := append([]storedSessionEvent{}, store.events...) + store.mu.Unlock() + for _, ep := range snapshot { + se, err := decodeSessionEvent[*schema.Message](ep.Data) + require.NoError(t, err) + if se.Message != nil && se.Message.Role == schema.Tool && se.Message.Content == "tool result" { + stored = true + } + } + assert.True(t, stored, "materialized tool-result stream must be persisted by finalization") +} + +func TestStreamPersistence_AgenticToolResultChunksConcat(t *testing.T) { + ctx := context.Background() + store := newAgenticSessionHelperStore() + sid := "agentic-tool-stream-session" + + agent := &agenticSessionStreamingAgent{ + chunks: []*schema.AgenticMessage{ + agenticToolResultMessage("call_1", "execute", "first\n"), + agenticToolResultMessage("call_1", "execute", "second\n"), + }, + turnEnd: &testTurnState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("q"), + agenticToolResultMessage("call_1", "execute", "first\nsecond\n"), + }, + }, + } + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("q")}) + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + for { + _, err := ev.Output.MessageOutput.MessageStream.Recv() + if err == io.EOF { + break + } + require.NoError(t, err) + } + } + } + + var stored *SessionEvent[*schema.AgenticMessage] + res, err := store.LoadEventsForSession(ctx, sid, nil) + require.NoError(t, err) + for _, se := range res.Events { + if se.Kind == SessionEventMessage && se.Message != nil && + len(se.Message.ContentBlocks) == 1 && + se.Message.ContentBlocks[0].Type == schema.ContentBlockTypeFunctionToolResult { + stored = se + break + } + } + + require.NotNil(t, stored) + require.NotNil(t, stored.Message) + require.Len(t, stored.Message.ContentBlocks, 1) + ftr := stored.Message.ContentBlocks[0].FunctionToolResult + require.NotNil(t, ftr) + assert.Equal(t, "call_1", ftr.CallID) + assert.Equal(t, "execute", ftr.Name) + require.Len(t, ftr.Content, 1) + assert.Equal(t, "first\nsecond\n", ftr.Content[0].Text.Text) + assert.Nil(t, stored.Message.ContentBlocks[0].StreamingMeta) +} + +func TestStreamPersistence_AgenticToolResultChunksWithStreamingMeta(t *testing.T) { + ctx := context.Background() + store := newAgenticSessionHelperStore() + sid := "agentic-tool-stream-meta-session" + + first := agenticToolResultMessage("call_1", "execute", "first\n") + second := agenticToolResultMessage("call_1", "execute", "second\n") + first.ContentBlocks[0].StreamingMeta = &schema.StreamingMeta{Index: 0} + second.ContentBlocks[0].StreamingMeta = &schema.StreamingMeta{Index: 0} + + agent := &agenticSessionStreamingAgent{ + chunks: []*schema.AgenticMessage{first, second}, + turnEnd: &testTurnState[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("q"), + agenticToolResultMessage("call_1", "execute", "first\nsecond\n"), + }, + }, + } + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("q")}) + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + for { + _, err := ev.Output.MessageOutput.MessageStream.Recv() + if err == io.EOF { + break + } + require.NoError(t, err) + } + } + } + + var stored *schema.AgenticMessage + res, err := store.LoadEventsForSession(ctx, sid, nil) + require.NoError(t, err) + for _, se := range res.Events { + if se.Kind == SessionEventMessage && se.Message != nil && + len(se.Message.ContentBlocks) == 1 && + se.Message.ContentBlocks[0].Type == schema.ContentBlockTypeFunctionToolResult { + stored = se.Message + break + } + } + + require.NotNil(t, stored) + require.Len(t, stored.ContentBlocks, 1) + block := stored.ContentBlocks[0] + assert.Nil(t, block.StreamingMeta) + require.NotNil(t, block.FunctionToolResult) + assert.Equal(t, "call_1", block.FunctionToolResult.CallID) + assert.Equal(t, "execute", block.FunctionToolResult.Name) + require.Len(t, block.FunctionToolResult.Content, 1) + assert.Equal(t, "first\nsecond\n", block.FunctionToolResult.Content[0].Text.Text) +} + +func TestStreamPersistence_AgenticIncompleteStreamPrefixPersisted(t *testing.T) { + ctx := context.Background() + store := newAgenticSessionHelperStore() + sid := "agentic-incomplete-stream-session" + streamErr := errors.New("agentic model stream failed") + chunk := agenticToolResultMessage("call_1", "execute", "partial\n") + agent := &agenticSessionStreamingAgent{ + chunks: []*schema.AgenticMessage{chunk}, + streamErr: streamErr, + } + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("q")}) + var sawStreamErr bool + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + for { + _, err := ev.Output.MessageOutput.MessageStream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + assert.ErrorContains(t, err, streamErr.Error()) + sawStreamErr = true + break + } + } + } + } + require.True(t, sawStreamErr) + + res, err := store.LoadEventsForSession(ctx, sid, nil) + require.NoError(t, err) + var incomplete []*SessionEvent[*schema.AgenticMessage] + var normalToolMessages []*SessionEvent[*schema.AgenticMessage] + for _, se := range res.Events { + if se.Kind == SessionEventMessageStreamIncomplete { + incomplete = append(incomplete, se) + } + if se.Kind == SessionEventMessage && se.Message != nil && + len(se.Message.ContentBlocks) == 1 && + se.Message.ContentBlocks[0].Type == schema.ContentBlockTypeFunctionToolResult { + normalToolMessages = append(normalToolMessages, se) + } + } + require.Len(t, incomplete, 1) + require.NotNil(t, incomplete[0].MessageStreamIncomplete) + prefix := incomplete[0].MessageStreamIncomplete.Message + require.NotNil(t, prefix) + require.Len(t, prefix.ContentBlocks, 1) + require.NotNil(t, prefix.ContentBlocks[0].FunctionToolResult) + require.Len(t, prefix.ContentBlocks[0].FunctionToolResult.Content, 1) + assert.Equal(t, "partial\n", prefix.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) + assert.Contains(t, incomplete[0].MessageStreamIncomplete.Error, streamErr.Error()) + assert.Empty(t, normalToolMessages) + + reconstructed, err := reconstructSessionState[*schema.AgenticMessage](ctx, mustOpenTestSession[*schema.AgenticMessage](t, ctx, store, sid), sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, reconstructed) + require.NotNil(t, reconstructed.state) + require.Len(t, reconstructed.state.Messages, 1) + assert.Equal(t, schema.AgenticRoleTypeUser, reconstructed.state.Messages[0].Role) +} + +func agenticToolResultMessage(callID, name, text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: callID, + Name: name, + Content: []*schema.FunctionToolResultContentBlock{ + { + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: text}, + }, + }, + }, + }, + }, + } +} + +// TestStreamPersistence_GetMessageError_NotEnqueued verifies that a stream +// materialization error sets persistErr (failing the turn commit) and does NOT +// enqueue a corrupt SessionEvent. +func TestStreamPersistence_GetMessageError_NotEnqueued(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "stream-err-session" + + // Build a stream that errors on Recv. + streamReader, streamWriter := schema.Pipe[*schema.Message](2) + streamWriter.Send(schema.AssistantMessage("partial ", nil), nil) + streamWriter.Send(nil, errors.New("simulated stream failure")) + streamWriter.Close() + + agent := &streamingAgentRaw{ + stream: streamReader, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("ok", nil)}, + }, + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + iter := runner.Query(ctx, "trigger") + var lastErr error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + lastErr = ev.Err + } + // Drain any live stream so the goroutine doesn't leak. + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + _, _ = schema.ConcatMessageStream(ev.Output.MessageOutput.MessageStream) + } + } + require.NoError(t, lastErr, "stream materialization errors should drop only the message event") + + // Verify no assistant SessionEvent is in the log. + for _, ep := range store.events { + se, err := decodeSessionEvent[*schema.Message](ep.Data) + require.NoError(t, err) + if se.Message != nil { + assert.NotEqual(t, schema.Assistant, se.Message.Role, + "failed stream must not produce a persisted assistant event") + } + } +} + +func TestStreamPersistence_GetMessageErrorSurfacesAfterLiveStreaming(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "sync-stream-err-session" + + streamReader, streamWriter := schema.Pipe[*schema.Message](2) + streamWriter.Send(schema.AssistantMessage("partial ", nil), nil) + streamWriter.Send(nil, errors.New("simulated stream failure")) + streamWriter.Close() + + agent := &streamingAgentRaw{ + stream: streamReader, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("ok", nil)}, + }, + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + SessionID: sid, + SessionStore: store, + }) + + iter := runner.Query(ctx, "trigger") + var lastErr error + var sawOutput bool + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + lastErr = ev.Err + } + if ev.Output != nil && ev.Output.MessageOutput != nil { + sawOutput = true + } + } + require.NoError(t, lastErr) + assert.True(t, sawOutput, "streaming output may already be live before materialization fails") + + for _, ep := range store.events { + se, err := decodeSessionEvent[*schema.Message](ep.Data) + require.NoError(t, err) + if se.Message != nil { + assert.NotEqual(t, schema.Assistant, se.Message.Role, + "failed sync stream must not produce a persisted assistant event") + } + } +} + +// streamingAgentRaw lets the test inject an arbitrary stream reader (including +// one that emits errors). +type streamingAgentRaw struct { + stream *schema.StreamReader[*schema.Message] + turnEnd *testTurnState[*schema.Message] +} + +func (a *streamingAgentRaw) Name(_ context.Context) string { return "streaming-raw" } +func (a *streamingAgentRaw) Description(_ context.Context) string { return "stream-error test agent" } +func (a *streamingAgentRaw) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + mv := &MessageVariant{IsStreaming: true, MessageStream: a.stream, Role: schema.Assistant} + gen.Send(&AgentEvent{AgentName: "streaming-raw", Output: &AgentOutput{MessageOutput: mv}}) + }() + return iter +} + +// TestSessionEvent_NilVsEmptyMessagesReplaced verifies that nil and empty +// MessagesReplaced are distinguishable after round-trip through the serializer. +func TestSessionEvent_NilVsEmptyMessagesReplaced(t *testing.T) { + t.Run("nil MessagesReplaced", func(t *testing.T) { + msg := schema.UserMessage("just a message") + EnsureMessageID(msg) + se := &SessionEvent[*schema.Message]{Message: msg} + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + assert.Nil(t, decoded.MessagesReplaced, "absent MessagesReplaced must decode as nil pointer") + require.NotNil(t, decoded.Message) + }) + + t.Run("empty MessagesReplaced", func(t *testing.T) { + empty := []*schema.Message{} + se := &SessionEvent[*schema.Message]{MessagesReplaced: &empty} + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.MessagesReplaced, "&[]M{} must decode as non-nil pointer") + assert.Empty(t, *decoded.MessagesReplaced) + }) +} + +// TestRunnerInputEvents_MixedRoles verifies that callers can pass system + user +// messages and both are persisted with their original roles. +func TestRunnerInputEvents_MixedRoles(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "mixed-roles" + + agent := &runnerSessionAgent{ + name: "mr-agent", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.AssistantMessage("ok", nil)}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + + systemMsg := schema.SystemMessage("system instruction") + userMsg := schema.UserMessage("hello") + drainSessionEvents(t, runner.Run(ctx, []*schema.Message{systemMsg, userMsg})) + + // Find the first two message events: they must be the input messages with + // preserved roles. Lifecycle timeline records may surround them. + messageEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage + }) + require.GreaterOrEqual(t, len(messageEvents), 2) + first := messageEvents[0] + require.NotNil(t, first.Message) + assert.Equal(t, schema.System, first.Message.Role) + assert.Equal(t, "system instruction", first.Message.Content) + + second := messageEvents[1] + require.NotNil(t, second.Message) + assert.Equal(t, schema.User, second.Message.Role) + assert.Equal(t, "hello", second.Message.Content) +} + +func TestCustomAgentNormalCloseCommitsIdle(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "turn-end-only" + + agent := &turnEndOnlyAgent{} + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "input")) + + var sawCommit bool + for _, ep := range store.events { + se, err := decodeSessionEvent[*schema.Message](ep.Data) + require.NoError(t, err) + if isCommittedIdleEvent(se) { + sawCommit = true + } + } + assert.True(t, sawCommit) +} + +type turnEndOnlyAgent struct{} + +func (a *turnEndOnlyAgent) Name(_ context.Context) string { return "turn-end-only" } +func (a *turnEndOnlyAgent) Description(_ context.Context) string { return "" } +func (a *turnEndOnlyAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + }() + return iter +} + +// TestTailReplay_PartialTurnWithoutCommittedIdle verifies that events appended +// after the last committed idle are replayed on reconstruction. +func TestTailReplay_PartialTurnWithoutCommittedIdle(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "tail-replay" + + // Phase 1: a normal completed turn (messages + committed idle event). + a1 := schema.UserMessage("Q1") + EnsureMessageID(a1) + r1 := schema.AssistantMessage("A1", nil) + EnsureMessageID(r1) + for _, m := range []*schema.Message{a1, r1} { + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + committedIdleSE := withTestCommittedIdle[*schema.Message]("turn-1") + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{committedIdleSE})) + + // Phase 2: simulate a partial second turn where events were appended but + // no committed idle was persisted (interrupted). + a2 := schema.UserMessage("Q2") + EnsureMessageID(a2) + r2 := schema.AssistantMessage("A2", nil) + EnsureMessageID(r2) + for _, m := range []*schema.Message{a2, r2} { + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + // Boot: prepareRunnerSessionRun reconstructs durable context through the log tail. + state, err := prepareRunnerSessionRun[*schema.Message](ctx, nil, nil, sid, store, nil) + require.NoError(t, err) + require.True(t, state.enabled) + require.Len(t, state.latestState.Messages, 4) + assert.Equal(t, "Q1", state.latestState.Messages[0].Content) + assert.Equal(t, "A1", state.latestState.Messages[1].Content) + assert.Equal(t, "Q2", state.latestState.Messages[2].Content) + assert.Equal(t, "A2", state.latestState.Messages[3].Content) +} + +// TestTailReplay_NoTailEvents verifies that the fast path is not disturbed when +// no events follow the snapshot. +func TestTailReplay_NoTailEvents(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "no-tail" + + q := schema.UserMessage("Q") + EnsureMessageID(q) + se := withTestEventID(&SessionEvent[*schema.Message]{Message: q}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + + turnEndSE := withTestCommittedIdle[*schema.Message]("turn-1") + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{turnEndSE})) + + state, err := prepareRunnerSessionRun[*schema.Message](ctx, nil, nil, sid, store, nil) + require.NoError(t, err) + require.Len(t, state.latestState.Messages, 1) + assert.Equal(t, "Q", state.latestState.Messages[0].Content) +} + +// TestTailReplay_EmptySnapshotCursor verifies cursor-based replay correctly +// handles a snapshot that committed an empty Messages array — the cursor still +// excludes pre-boundary events. +func TestTailReplay_EmptySnapshotCursor(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "empty-snapshot" + + // Pre-boundary events. + for i := 0; i < 3; i++ { + m := schema.UserMessage("pre") + EnsureMessageID(m) + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + // MessagesReplaced boundary with empty slice — supersedes pre-boundary events. + empty := []*schema.Message{} + boundarySE := withTestEventID(&SessionEvent[*schema.Message]{MessagesReplaced: &empty}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{boundarySE})) + + // Post-boundary events. + postMsg := schema.UserMessage("post") + EnsureMessageID(postMsg) + se := withTestEventID(&SessionEvent[*schema.Message]{Message: postMsg}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + + state, err := prepareRunnerSessionRun[*schema.Message](ctx, nil, nil, sid, store, nil) + require.NoError(t, err) + require.Len(t, state.latestState.Messages, 1) + assert.Equal(t, "post", state.latestState.Messages[0].Content) +} + +type agenticSessionHelperStore struct { + mu sync.Mutex + events []storedSessionEvent + eventIDIdx map[string]int +} + +func newAgenticSessionHelperStore() *agenticSessionHelperStore { + return &agenticSessionHelperStore{eventIDIdx: make(map[string]int)} +} + +func (s *agenticSessionHelperStore) AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[*schema.AgenticMessage]) error { + return s.AppendEventsForSession(ctx, sessionID, events) +} + +func (s *agenticSessionHelperStore) AppendEventsForSession(_ context.Context, _ string, events []*SessionEvent[*schema.AgenticMessage]) error { + s.mu.Lock() + defer s.mu.Unlock() + for _, event := range events { + if event == nil || event.EventID == "" { + return ErrInvalidEventID + } + if err := NormalizeSessionEventKind(event); err != nil { + return err + } + if _, ok := s.eventIDIdx[event.EventID]; ok { + continue + } + data, err := encodeSessionEvent(event) + if err != nil { + return err + } + s.events = append(s.events, storedSessionEvent{EventID: event.EventID, Kind: event.Kind, Data: data}) + s.eventIDIdx[event.EventID] = len(s.events) - 1 + } + return nil +} + +func (s *agenticSessionHelperStore) LoadEvents(ctx context.Context, sessionID string, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.AgenticMessage], error) { + return s.LoadEventsForSession(ctx, sessionID, req) +} + +func (s *agenticSessionHelperStore) LoadEventsForSession(_ context.Context, _ string, opts *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.AgenticMessage], error) { + s.mu.Lock() + defer s.mu.Unlock() + if opts == nil { + opts = &LoadSessionEventsRequest{} + } + start, end, step := 0, len(s.events), 1 + if opts.After != "" { + pos, ok := s.eventIDIdx[opts.After] + if !ok { + return nil, ErrEventIDOutOfRange + } + if opts.Reverse { + start, end, step = pos-1, -1, -1 + } else { + start = pos + 1 + } + } else if opts.Reverse { + start, end, step = len(s.events)-1, -1, -1 + } + kindSet := buildTestKindSet(opts.Kinds) + var out []*SessionEvent[*schema.AgenticMessage] + for i := start; i != end; i += step { + if i < 0 || i >= len(s.events) { + break + } + rec := s.events[i] + if kindSet != nil { + if _, ok := kindSet[rec.Kind]; !ok { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + break + } + event, err := decodeSessionEvent[*schema.AgenticMessage](rec.Data) + if err != nil { + return nil, err + } + out = append(out, event) + } + return &LoadSessionEventsResult[*schema.AgenticMessage]{Events: out}, nil +} + +func (s *agenticSessionHelperStore) openSession(_ context.Context, req *openSessionRequest) (*openSessionResult[*schema.AgenticMessage], error) { + sessionID := "" + if req != nil { + sessionID = req.sessionID + } + return &openSessionResult[*schema.AgenticMessage]{ + handle: &agenticTestSessionHandle{store: s, sessionID: sessionID}, + }, nil +} + +type agenticTestSessionHandle struct { + store *agenticSessionHelperStore + sessionID string +} + +func (h *agenticTestSessionHandle) loadEvents(ctx context.Context, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.AgenticMessage], error) { + if req == nil { + req = &LoadSessionEventsRequest{} + } + return h.store.LoadEventsForSession(ctx, h.sessionID, req) +} + +func (h *agenticTestSessionHandle) appendEvents(ctx context.Context, events []*SessionEvent[*schema.AgenticMessage]) error { + return h.store.AppendEventsForSession(ctx, h.sessionID, events) +} + +func (h *agenticTestSessionHandle) close(context.Context) error { return nil } + +// TestPartialInterrupted_ThenNewRun verifies that when a turn is interrupted +// after some events have been appended (but before the committed idle marker), a new +// Run with NO CheckPointStore (i.e. session-only mode) recovers the in-flight +// events via tail replay rather than treating the session as fresh. +// +// This test does not use CheckPointStore — Runner skips pending checkpoints +// on fresh Run, so checkpoint presence would not block regardless. +func TestPartialInterrupted_ThenNewRun(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "partial-interrupted" + + // Phase 1: simulate a normal completed turn. + q1 := schema.UserMessage("first") + EnsureMessageID(q1) + r1 := schema.AssistantMessage("answer1", nil) + EnsureMessageID(r1) + for _, m := range []*schema.Message{q1, r1} { + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + committedIdleSE := withTestCommittedIdle[*schema.Message]("turn-1") + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{committedIdleSE})) + + // Phase 2: simulate an interrupted turn with events appended but no committed idle. + q2 := schema.UserMessage("partial") + EnsureMessageID(q2) + for _, m := range []*schema.Message{q2} { + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + // Phase 3: new Run (no CheckPointStore; Runner skips pending checkpoints on fresh Run). + captured := &runnerSessionAgent{ + name: "ra", + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: captured, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "second")) + + // Fresh Run includes durable partial-turn context because Session + // reconstruction replays context events through the log tail. + require.Len(t, captured.inputs, 1) + contents := []string{} + for _, m := range captured.inputs[0] { + contents = append(contents, m.Content) + } + assert.Equal(t, []string{"first", "answer1", "partial", "second"}, contents) +} + +// TestSessionEvent_StreamCopyConcat_ByteIdentical verifies the round-trip of a +// streamed-then-persisted SessionEvent matches what the live consumer sees. +func TestSessionEvent_StreamCopyConcat_ByteIdentical(t *testing.T) { + chunks := []*schema.Message{ + schema.AssistantMessage("foo ", nil), + schema.AssistantMessage("bar ", nil), + schema.AssistantMessage("baz", nil), + } + stream := schema.StreamReaderFromArray(chunks) + + // Mimic the runner's logic: copy, materialize one side, leave the other live. + copies := stream.Copy(2) + persistCopy := &TypedMessageVariant[*schema.Message]{IsStreaming: true, MessageStream: copies[0]} + persistedMsg, err := persistCopy.GetMessage() + require.NoError(t, err) + require.NotNil(t, persistedMsg) + + se := &SessionEvent[*schema.Message]{Message: persistedMsg} + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.Message) + assert.Equal(t, "foo bar baz", decoded.Message.Content) + + // The live copy should yield the same concatenated content. + liveMsg, err := schema.ConcatMessageStream(copies[1]) + require.NoError(t, err) + assert.Equal(t, decoded.Message.Content, liveMsg.Content) +} + +// TestExplicitCheckpointResume_WithSessionMode verifies that when a caller passes +// an explicit checkpoint ID alongside a configured SessionID/SessionStore[*schema.Message], the +// resume path still loads reconstructed session state. +func TestExplicitCheckpointResume_WithSessionMode(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "explicit-cp-session" + + // Seed the session store with events and a committed idle marker. + prior := &testTurnState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("seed"), schema.AssistantMessage("seed-ans", nil)}, + } + // Seed session events (messages + committed idle). + for _, m := range prior.Messages { + EnsureMessageID(m) + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + committedIdleSE := withTestCommittedIdle[*schema.Message]("turn-1") + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{committedIdleSE})) + + // Seed an arbitrary checkpoint ID with a runner-session-checkpoint wrapper + // so runnerLoadCheckPointForSession can decode it. + cpBytes, err := encodeRunnerSessionCheckpoint(&runnerSessionCheckpoint{ + Payload: []byte("opaque"), + }) + require.NoError(t, err) + explicitCheckpointID := "user-supplied-cp" + require.NoError(t, store.Set(ctx, explicitCheckpointID, cpBytes)) + + state, effective, err := prepareRunnerSessionResume[*schema.Message](ctx, store, sid, store, nil, explicitCheckpointID) + require.NoError(t, err) + require.True(t, state.enabled, "session mode must remain enabled when an explicit checkpoint ID is supplied") + require.NotNil(t, state.latestState) + assert.Equal(t, 2, len(state.latestState.Messages), + "latest snapshot must be loaded for explicit-checkpoint resume in session mode") + assert.Equal(t, explicitCheckpointID, effective, + "caller-supplied checkpoint ID must be preserved") +} + +// TestResumePath_TailReplay verifies that the resume path also performs tail +// replay (uses the same fast path as the run path). +func TestResumePath_TailReplay(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "resume-tail" + + q1 := schema.UserMessage("Q") + EnsureMessageID(q1) + r1 := schema.AssistantMessage("A", nil) + EnsureMessageID(r1) + for _, m := range []*schema.Message{q1, r1} { + se := withTestEventID(&SessionEvent[*schema.Message]{Message: m}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + turnEndSE := withTestCommittedIdle[*schema.Message]("turn-1") + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{turnEndSE})) + + // Append a tail event after the snapshot. + tailMsg := schema.UserMessage("post-snapshot") + EnsureMessageID(tailMsg) + se := withTestEventID(&SessionEvent[*schema.Message]{Message: tailMsg}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + + // Seed a runner session checkpoint so the resume path finds something to load. + cpStore := newSessionHelperStore() + cpBytes, err := encodeRunnerSessionCheckpoint(&runnerSessionCheckpoint{ + Payload: []byte("opaque"), + }) + require.NoError(t, err) + require.NoError(t, cpStore.Set(ctx, sessionRunnerCheckpointID(sid), cpBytes)) + + state, _, err := prepareRunnerSessionResume[*schema.Message](ctx, cpStore, sid, store, nil, "") + require.NoError(t, err) + require.Len(t, state.latestState.Messages, 3, + "resume boot state should include durable context events through the log tail") + assert.Equal(t, "Q", state.latestState.Messages[0].Content) + assert.Equal(t, "A", state.latestState.Messages[1].Content) + assert.Equal(t, "post-snapshot", state.latestState.Messages[2].Content) +} + +// Ensure the io package import is used (for compile when chunks are empty). + +// mutationAgent emits a sequence of caller-provided TypedAgentEvents. Used to +// verify the runner persists each session-mutation +// event variant (MessagesReplaced, MessageUpdated, MessageInserted) faithfully. +type mutationAgent struct { + events []*AgentEvent + turnEnd *testTurnState[*schema.Message] +} + +func (a *mutationAgent) Name(_ context.Context) string { return "mutation-agent" } +func (a *mutationAgent) Description(_ context.Context) string { return "" } +func (a *mutationAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + for _, ev := range a.events { + gen.Send(ev) + } + }() + return iter +} + +// TestRunnerPersists_MessagesReplaced verifies a MessagesReplaced event from +// any source (e.g. summarization) is persisted. +func TestRunnerPersists_MessagesReplaced(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "mr-session" + + summary := schema.AssistantMessage("summary content", nil) + EnsureMessageID(summary) + repl := []*schema.Message{summary} + + agent := &mutationAgent{ + events: []*AgentEvent{ + { + AgentName: "mutation-agent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessagesReplaced, + MessagesReplaced: &repl, + }, + }, + }, + }, + turnEnd: &testTurnState[*schema.Message]{Messages: []*schema.Message{summary}}, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "anything")) + + // Read events back via the store. + res, err := store.LoadEventsForSession(ctx, sid, &LoadSessionEventsRequest{}) + require.NoError(t, err) + + var foundReplaced bool + for _, se := range res.Events { + if se.MessagesReplaced != nil { + foundReplaced = true + require.Len(t, *se.MessagesReplaced, 1) + assert.Equal(t, "summary content", (*se.MessagesReplaced)[0].Content) + } + } + assert.True(t, foundReplaced, "MessagesReplaced must be persisted") +} + +// TestRunnerPersists_MessageUpdated_BothMessages verifies that when reduction +// emits two MessageUpdated events (one for the assistant tool-call message, +// one for the tool-result message), both reach the event log and reconstruction +// applies them correctly. +func TestRunnerPersists_MessageUpdated_BothMessages(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "mu-session" + + // Build two messages with stable IDs. + toolCallMsg := schema.AssistantMessage("call me", nil) + EnsureMessageID(toolCallMsg) + toolResultMsg := schema.ToolMessage("result content", "tc-1", schema.WithToolName("t1")) + EnsureMessageID(toolResultMsg) + + // Pretend reduction rewrites both: the assistant message's args (we just + // reuse the same message pointer for the test, with a marker) and the tool + // result content. + updatedAssistant := schema.AssistantMessage("call me [cleared]", nil) + updatedAssistant.Extra = map[string]any{"_eino_msg_id": GetMessageID(toolCallMsg), "cleared": true} + updatedTool := schema.ToolMessage("[placeholder]", "tc-1", schema.WithToolName("t1")) + updatedTool.Extra = map[string]any{"_eino_msg_id": GetMessageID(toolResultMsg)} + + agent := &mutationAgent{ + events: []*AgentEvent{ + { + AgentName: "mutation-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: toolCallMsg, Role: schema.Assistant}, + }, + }, + { + AgentName: "mutation-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: toolResultMsg, Role: schema.Tool, ToolName: "t1"}, + }, + }, + { + AgentName: "mutation-agent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessageUpdated, + MessageUpdated: &MessageUpdatedEvent[*schema.Message]{ + MessageID: GetMessageID(toolResultMsg), + Message: updatedTool, + }, + }, + }, + }, + { + AgentName: "mutation-agent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessageUpdated, + MessageUpdated: &MessageUpdatedEvent[*schema.Message]{ + MessageID: GetMessageID(toolCallMsg), + Message: updatedAssistant, + }, + }, + }, + }, + }, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{updatedAssistant, updatedTool}, + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Query(ctx, "go")) + + res, err := store.LoadEventsForSession(ctx, sid, &LoadSessionEventsRequest{}) + require.NoError(t, err) + + var updates int + for _, se := range res.Events { + if se.MessageUpdated != nil { + updates++ + } + } + assert.Equal(t, 2, updates, "both MessageUpdated events must be persisted") + + // Reconstruction must apply both updates correctly. + result, err := reconstructSessionState[*schema.Message](ctx, mustOpenTestSession[*schema.Message](t, ctx, store, sid), sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + // Find updated content among reconstructed messages. + var sawClearedAssistant, sawPlaceholderTool bool + for _, m := range result.state.Messages { + if m.Role == schema.Assistant && m.Content == "call me [cleared]" { + sawClearedAssistant = true + } + if m.Role == schema.Tool && m.Content == "[placeholder]" { + sawPlaceholderTool = true + } + } + assert.True(t, sawClearedAssistant, "reconstruction must apply cleared assistant update") + assert.True(t, sawPlaceholderTool, "reconstruction must apply placeholder tool update") +} + +// TestRunnerPersists_MessageInserted_AnchorAndAppend verifies that +// MessageInserted events from middlewares (AgentsMD, ToolSearch, PatchToolCalls) +// flow through the runner, are persisted, and reconstruct correctly. +func TestRunnerPersists_MessageInserted_AnchorAndAppend(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "mi-session" + + // Anchor: the user message in the session, present from the input. + userMsg := schema.UserMessage("hello") + EnsureMessageID(userMsg) + + // AgentsMD-style insertion before the user message. + agentsmdMsg := schema.UserMessage("[agentsmd content]") + agentsmdMsg.Extra = map[string]any{"__agentsmd_content__": true} + EnsureMessageID(agentsmdMsg) + + // PatchToolCalls-style append at end. + patchedTool := schema.ToolMessage("[patched]", "tc-1", schema.WithToolName("t1")) + EnsureMessageID(patchedTool) + + finalMessages := []*schema.Message{agentsmdMsg, userMsg, patchedTool} + + agent := &mutationAgent{ + events: []*AgentEvent{ + // Mimic input event flow: user message already appears in the input. + // MessageInserted before the user message: + { + AgentName: "mutation-agent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessageInserted, + MessageInserted: &MessageInsertedEvent[*schema.Message]{ + Message: agentsmdMsg, + BeforeMessageID: GetMessageID(userMsg), + }, + }, + }, + }, + // MessageInserted appended at end: + { + AgentName: "mutation-agent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessageInserted, + MessageInserted: &MessageInsertedEvent[*schema.Message]{ + Message: patchedTool, + BeforeMessageID: "", + }, + }, + }, + }, + }, + turnEnd: &testTurnState[*schema.Message]{Messages: finalMessages}, + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + // We must pass the user message as input, with its existing ID already assigned, + // so reconstruction's anchor lookup succeeds. + drainSessionEvents(t, runner.Run(ctx, []*schema.Message{userMsg})) + + res, err := store.LoadEventsForSession(ctx, sid, &LoadSessionEventsRequest{}) + require.NoError(t, err) + + var inserts int + for _, se := range res.Events { + if se.MessageInserted != nil { + inserts++ + } + } + assert.Equal(t, 2, inserts, "both MessageInserted events must be persisted") + + // Verify reconstruction applies insertions correctly. + result, err := reconstructSessionState[*schema.Message](ctx, mustOpenTestSession[*schema.Message](t, ctx, store, sid), sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.GreaterOrEqual(t, len(result.state.Messages), 3) + // The agentsmd message should appear before the user input. + var idxAgentsmd, idxUser, idxPatched int + idxAgentsmd, idxUser, idxPatched = -1, -1, -1 + for i, m := range result.state.Messages { + switch GetMessageID(m) { + case GetMessageID(agentsmdMsg): + idxAgentsmd = i + case GetMessageID(userMsg): + idxUser = i + case GetMessageID(patchedTool): + idxPatched = i + } + } + require.NotEqual(t, -1, idxAgentsmd) + require.NotEqual(t, -1, idxUser) + require.NotEqual(t, -1, idxPatched) + assert.Less(t, idxAgentsmd, idxUser, "agentsmd must be inserted before the user message") + assert.Greater(t, idxPatched, idxUser, "patched tool message must be appended at the end") +} + +type leadingSystemTestModel[M MessageType] struct { + response M + inputs [][]M +} + +func (m *leadingSystemTestModel[M]) Generate(_ context.Context, input []M, _ ...model.Option) (M, error) { + copied := append([]M{}, input...) + m.inputs = append(m.inputs, copied) + return m.response, nil +} + +func (m *leadingSystemTestModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]M{msg}), nil +} + +type sessionToolCallingModel struct { + response *schema.Message + inputs [][]*schema.Message +} + +func (m *sessionToolCallingModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.inputs = append(m.inputs, append([]*schema.Message{}, input...)) + return m.response, nil +} + +func (m *sessionToolCallingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *sessionToolCallingModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +type modelContextExtraTool struct{} + +func (modelContextExtraTool) Info(context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "extra_tool", + Desc: "tool with json-normalized extra metadata", + Extra: map[string]any{"version": 1}, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": { + Type: schema.String, + Desc: "query", + }, + }), + }, nil +} + +func (modelContextExtraTool) InvokableRun(context.Context, string, ...tool.Option) (string, error) { + return "ok", nil +} + +func drainAgenticSessionEvents(t *testing.T, iter *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) { + t.Helper() + for { + event, ok := iter.Next() + if !ok { + return + } + require.NoError(t, event.Err) + } +} + +func loadMessageSessionEvents(t *testing.T, ctx context.Context, store *sessionHelperStore, sid string) []*SessionEvent[*schema.Message] { + t.Helper() + res, err := store.LoadEventsForSession(ctx, sid, &LoadSessionEventsRequest{}) + require.NoError(t, err) + return res.Events +} + +func loadAgenticSessionEvents(t *testing.T, ctx context.Context, store *agenticSessionHelperStore, sid string) []*SessionEvent[*schema.AgenticMessage] { + t.Helper() + res, err := store.LoadEventsForSession(ctx, sid, &LoadSessionEventsRequest{}) + require.NoError(t, err) + return res.Events +} + +func TestRunnerNoPersist_GeneratedLeadingSystemMessage(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "no-persist-gen-system" + model := &leadingSystemTestModel[*schema.Message]{response: schema.AssistantMessage("answer", nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "gen-system-agent", + Description: "test", + Instruction: "system v1", + Model: model, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + drainSessionEvents(t, runner.Run(ctx, []*schema.Message{schema.UserMessage("hello")})) + + for _, event := range loadMessageSessionEvents(t, ctx, store, sid) { + if event.Message != nil && event.Message.Role == schema.System { + t.Fatalf("generated leading system message must not be persisted as SessionEventMessage") + } + if event.MessageInserted != nil && event.MessageInserted.Message.Role == schema.System { + t.Fatalf("generated leading system message must not be persisted as SessionEventMessageInserted") + } + if event.MessageUpdated != nil && event.MessageUpdated.Message.Role == schema.System { + t.Fatalf("generated leading system message must not be persisted as SessionEventMessageUpdated") + } + } + + require.Len(t, model.inputs, 1) + require.GreaterOrEqual(t, len(model.inputs[0]), 2) + assert.Equal(t, schema.System, model.inputs[0][0].Role) + assert.Equal(t, "system v1", model.inputs[0][0].Content) + + handle := mustOpenTestSession[*schema.Message](t, ctx, store, sid) + result, err := reconstructSessionState[*schema.Message](ctx, handle, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NoError(t, handle.close(ctx)) + for _, msg := range result.state.Messages { + if msg.Role == schema.System { + t.Fatalf("reconstructed state must not contain generated leading system message") + } + } +} + +func TestRunnerNoPersist_GeneratedLeadingSystemEmptySession(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "no-persist-gen-system-empty" + model := &leadingSystemTestModel[*schema.Message]{response: schema.AssistantMessage("answer", nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "gen-system-empty-agent", + Description: "test", + Instruction: "system only", + Model: model, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + drainSessionEvents(t, runner.Run(ctx, nil)) + + for _, event := range loadMessageSessionEvents(t, ctx, store, sid) { + if event.Message != nil && event.Message.Role == schema.System { + t.Fatalf("generated leading system message must not be persisted in empty session") + } + } + require.Len(t, model.inputs, 1) + require.GreaterOrEqual(t, len(model.inputs[0]), 1) + assert.Equal(t, schema.System, model.inputs[0][0].Role) + assert.Equal(t, "system only", model.inputs[0][0].Content) +} + +func TestRunnerRecalculatesSystemMessageOnSecondRun(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "recalc-system-second-run" + + runTurn := func(instruction, user string) *leadingSystemTestModel[*schema.Message] { + model := &leadingSystemTestModel[*schema.Message]{response: schema.AssistantMessage("answer "+user, nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "recalc-system-agent", + Description: "test", + Instruction: instruction, + Model: model, + }) + require.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + drainSessionEvents(t, runner.Run(ctx, []*schema.Message{schema.UserMessage(user)})) + return model + } + + m1 := runTurn("system v1", "one") + require.Len(t, m1.inputs, 1) + require.GreaterOrEqual(t, len(m1.inputs[0]), 2) + assert.Equal(t, "system v1", m1.inputs[0][0].Content) + assert.Equal(t, schema.System, m1.inputs[0][0].Role) + + m2 := runTurn("system v2", "two") + require.Len(t, m2.inputs, 1) + require.GreaterOrEqual(t, len(m2.inputs[0]), 3) + assert.Equal(t, "system v2", m2.inputs[0][0].Content) + assert.Equal(t, schema.System, m2.inputs[0][0].Role) + assert.Equal(t, "one", m2.inputs[0][1].Content) +} + +func TestRunnerNoPersist_CustomGenModelInputLeadingSystem(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "no-persist-custom-gen-system" + model := &leadingSystemTestModel[*schema.Message]{response: schema.AssistantMessage("answer", nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "custom-gen-system-agent", + Description: "test", + Instruction: "ignored", + Model: model, + GenModelInput: func(_ context.Context, _ string, input *AgentInput) ([]*schema.Message, error) { + system := schema.SystemMessage("custom system") + messages := make([]*schema.Message, 0, len(input.Messages)+1) + messages = append(messages, system) + messages = append(messages, input.Messages...) + return messages, nil + }, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + drainSessionEvents(t, runner.Run(ctx, []*schema.Message{schema.UserMessage("hello")})) + + for _, event := range loadMessageSessionEvents(t, ctx, store, sid) { + if event.Message != nil && event.Message.Role == schema.System { + t.Fatalf("custom GenModelInput leading system must not be persisted") + } + if event.MessageInserted != nil && event.MessageInserted.Message.Role == schema.System { + t.Fatalf("custom GenModelInput leading system must not be inserted") + } + if event.MessageUpdated != nil && event.MessageUpdated.Message.Role == schema.System { + t.Fatalf("custom GenModelInput leading system must not be updated") + } + } + require.Len(t, model.inputs, 1) + require.GreaterOrEqual(t, len(model.inputs[0]), 2) + assert.Equal(t, schema.System, model.inputs[0][0].Role) + assert.Equal(t, "custom system", model.inputs[0][0].Content) +} + +func TestRunnerPreserves_CallerSuppliedLeadingSystemMessage(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "preserve-caller-system" + model := &leadingSystemTestModel[*schema.Message]{response: schema.AssistantMessage("answer", nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "caller-system-agent", + Description: "test", + Instruction: "", + Model: model, + }) + require.NoError(t, err) + + systemMsg := schema.SystemMessage("caller system prompt") + EnsureMessageID(systemMsg) + userMsg := schema.UserMessage("hello") + EnsureMessageID(userMsg) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + drainSessionEvents(t, runner.Run(ctx, []*schema.Message{systemMsg, userMsg})) + + var systemEvents int + for _, event := range loadMessageSessionEvents(t, ctx, store, sid) { + if event.Message != nil && event.Message.Role == schema.System { + systemEvents++ + } + } + assert.Equal(t, 1, systemEvents, "caller-supplied system message must be persisted") + + handle := mustOpenTestSession[*schema.Message](t, ctx, store, sid) + result, err := reconstructSessionState[*schema.Message](ctx, handle, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NoError(t, handle.close(ctx)) + require.GreaterOrEqual(t, len(result.state.Messages), 2) + assert.Equal(t, schema.System, result.state.Messages[0].Role) + assert.Equal(t, "caller system prompt", result.state.Messages[0].Content) +} + +func TestRunnerSkipsLeadingSystemEventWhenCustomGenModelInputHasNoSystem(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "leading-system-custom-none" + model := &leadingSystemTestModel[*schema.Message]{response: schema.AssistantMessage("answer", nil)} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "custom-no-system-agent", + Description: "test", + Instruction: "ignored by custom input", + Model: model, + GenModelInput: func(_ context.Context, _ string, input *AgentInput) ([]*schema.Message, error) { + return append([]*schema.Message{}, input.Messages...), nil + }, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: sid, SessionStore: store}) + drainSessionEvents(t, runner.Run(ctx, []*schema.Message{schema.UserMessage("hello")})) + + for _, event := range loadMessageSessionEvents(t, ctx, store, sid) { + switch { + case event.Message != nil && event.Message.Role == schema.System: + t.Fatalf("custom GenModelInput without leading system must not persist system message") + case event.MessageInserted != nil && event.MessageInserted.Message.Role == schema.System: + t.Fatalf("custom GenModelInput without leading system must not insert system message") + case event.MessageUpdated != nil && event.MessageUpdated.Message.Role == schema.System: + t.Fatalf("custom GenModelInput without leading system must not update system message") + } + } + require.Len(t, model.inputs, 1) + require.Len(t, model.inputs[0], 1) + assert.Equal(t, schema.User, model.inputs[0][0].Role) +} + +func TestRunnerNoPersist_AgenticLeadingSystemMessage(t *testing.T) { + ctx := context.Background() + store := newAgenticSessionHelperStore() + sid := "no-persist-agentic-system" + + runTurn := func(instruction, user string) *leadingSystemTestModel[*schema.AgenticMessage] { + model := &leadingSystemTestModel[*schema.AgenticMessage]{response: agenticAssistantMessage("answer " + user)} + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "agentic-system-agent", + Description: "test", + Instruction: instruction, + Model: model, + }) + require.NoError(t, err) + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + drainAgenticSessionEvents(t, runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage(user)})) + return model + } + + m1 := runTurn("agentic system v1", "one") + for _, event := range loadAgenticSessionEvents(t, ctx, store, sid) { + if event.Message != nil && event.Message.Role == schema.AgenticRoleTypeSystem { + t.Fatalf("agentic generated leading system must not be persisted as message") + } + if event.MessageInserted != nil && event.MessageInserted.Message.Role == schema.AgenticRoleTypeSystem { + t.Fatalf("agentic generated leading system must not be persisted as inserted") + } + if event.MessageUpdated != nil && event.MessageUpdated.Message.Role == schema.AgenticRoleTypeSystem { + t.Fatalf("agentic generated leading system must not be persisted as updated") + } + } + require.Len(t, m1.inputs, 1) + require.GreaterOrEqual(t, len(m1.inputs[0]), 2) + assert.Equal(t, schema.AgenticRoleTypeSystem, m1.inputs[0][0].Role) + + m2 := runTurn("agentic system v2", "two") + require.Len(t, m2.inputs, 1) + require.GreaterOrEqual(t, len(m2.inputs[0]), 3) + assert.Equal(t, schema.AgenticRoleTypeSystem, m2.inputs[0][0].Role) +} + +func TestRunnerPersists_MessagesDeleted_Reconstructs(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "md-session" + + a := schema.UserMessage("a") + b := schema.AssistantMessage("b", nil) + c := schema.UserMessage("c") + for _, msg := range []*schema.Message{a, b, c} { + EnsureMessageID(msg) + } + + agent := &mutationAgent{ + events: []*AgentEvent{ + { + AgentName: "mutation-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: a, Role: schema.User}, + }, + }, + { + AgentName: "mutation-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: b, Role: schema.Assistant}, + }, + }, + { + AgentName: "mutation-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: c, Role: schema.User}, + }, + }, + { + AgentName: "mutation-agent", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessagesDeleted, + MessagesDeleted: &MessagesDeletedEvent{ + MessageIDs: []string{GetMessageID(b)}, + }, + }, + }, + }, + }, + turnEnd: &testTurnState[*schema.Message]{Messages: []*schema.Message{a, c}}, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: store, + }) + drainSessionEvents(t, runner.Run(ctx, nil)) + + res, err := store.LoadEventsForSession(ctx, sid, &LoadSessionEventsRequest{}) + require.NoError(t, err) + + var foundDeleted bool + for _, se := range res.Events { + if se.MessagesDeleted != nil { + foundDeleted = true + assert.Equal(t, []string{GetMessageID(b)}, se.MessagesDeleted.MessageIDs) + } + } + assert.True(t, foundDeleted, "MessagesDeleted must be persisted") + + result, err := reconstructSessionState[*schema.Message](ctx, mustOpenTestSession[*schema.Message](t, ctx, store, sid), sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, result.state.Messages, 2) + assert.Equal(t, "a", result.state.Messages[0].Content) + assert.Equal(t, "c", result.state.Messages[1].Content) +} + +func TestReconstructSessionState_MessagesDeletedMissingTargetFails(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "md-missing-target" + + a := schema.UserMessage("a") + EnsureMessageID(a) + msgEvent := withTestEventID(&SessionEvent[*schema.Message]{Message: a}) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{msgEvent})) + + deleteEvent := withTestEventID(&SessionEvent[*schema.Message]{ + MessagesDeleted: &MessagesDeletedEvent{MessageIDs: []string{"ghost-id"}}, + }) + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{deleteEvent})) + + committedIdleEvent := withTestCommittedIdle[*schema.Message]("turn-1") + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{committedIdleEvent})) + + _, err := reconstructSessionState[*schema.Message](ctx, mustOpenTestSession[*schema.Message](t, ctx, store, sid), sid, defaultLoadPageSize) + require.Error(t, err) + assert.Contains(t, err.Error(), "ghost-id") +} + +// TestAgentTool_ChildSessionID_FiltersFromParentLog verifies that events +// forwarded from an inner agent (via AgentTool) are tagged with the child +// SessionEvent.SessionID and are NOT persisted into the parent's session event +// log. The parent's log only contains events that belong to its own session. +func TestAgentTool_ChildSessionID_FiltersFromParentLog(t *testing.T) { + ctx := context.Background() + parentStore := newSessionHelperStore() + sid := "parent-session" + + // Inner-agent forwarded event from AgentTool path. Tagging with a + // SessionEvent.SessionID that does not match the parent session must be + // filtered out of persistence. + childMsg := schema.AssistantMessage("inner-agent-output", nil) + EnsureMessageID(childMsg) + parentMsg := schema.AssistantMessage("parent-output", nil) + EnsureMessageID(parentMsg) + + agent := &mutationAgent{ + events: []*AgentEvent{ + // An event tagged as belonging to a different session — should not be persisted. + { + AgentName: "child", + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + SessionID: "agent_tool:abc-123", + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Message: childMsg, + }, + }, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: childMsg, Role: schema.Assistant}, + }, + }, + // The parent's own event — should be persisted. + { + AgentName: "parent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: parentMsg, Role: schema.Assistant}, + }, + }, + }, + turnEnd: &testTurnState[*schema.Message]{ + Messages: []*schema.Message{parentMsg}, + }, + } + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: sid, + SessionStore: parentStore, + }) + drainSessionEvents(t, runner.Query(ctx, "go")) + + // Verify that childMsg is NOT in the parent's persistent log, but parentMsg is. + res, err := parentStore.LoadEventsForSession(ctx, sid, &LoadSessionEventsRequest{}) + require.NoError(t, err) + var sawChild, sawParent bool + for _, se := range res.Events { + if se.Message != nil { + if GetMessageID(se.Message) == GetMessageID(childMsg) { + sawChild = true + } + if GetMessageID(se.Message) == GetMessageID(parentMsg) { + sawParent = true + } + } + } + assert.False(t, sawChild, "events tagged with a different SessionEvent.SessionID must NOT enter the parent session log") + assert.True(t, sawParent, "parent's own events must be persisted") +} + +// TestAgentToolInterruptState_RoundTrip verifies the wrapper struct round-trips +// through JSON and preserves the child SessionID for resume. +func TestAgentToolInterruptState_RoundTrip(t *testing.T) { + bridge := []byte("opaque-checkpoint-bytes") + wrapped := agentToolInterruptState{ + ChildSessionID: "agent_tool:abcd", + BridgeCheckpoint: bridge, + } + // Use the same JSON marshal/unmarshal path as agent_tool.go. + encoded, err := json.Marshal(wrapped) + require.NoError(t, err) + + var decoded agentToolInterruptState + require.NoError(t, json.Unmarshal(encoded, &decoded)) + assert.Equal(t, wrapped.ChildSessionID, decoded.ChildSessionID) + assert.Equal(t, wrapped.BridgeCheckpoint, decoded.BridgeCheckpoint) +} + +func TestAttack_SessionEventVariantBothSet(t *testing.T) { + ev := &TypedAgentEvent[Message]{ + SessionEventVariant: &SessionEventVariant[Message]{ + Event: &SessionEvent[Message]{ + EventID: "evt-1", + Kind: SessionEventMessage, + Message: schema.UserMessage("hello"), + }, + MessageStreamRef: &MessageStreamRef{ + EventID: "evt-2", + Kind: SessionEventMessage, + }, + }, + } + + err := validateAgentSessionEventIdentity(ev) + if err == nil { + t.Fatal("expected error when both Event and MessageStreamRef are set, got nil") + } + t.Logf("correctly rejected both-set variant: %v", err) +} + +func TestAttack_SessionEventVariantNeitherSet(t *testing.T) { + ev := &TypedAgentEvent[Message]{ + SessionEventVariant: &SessionEventVariant[Message]{ + SessionID: "sess-1", + }, + } + + err := validateAgentSessionEventIdentity(ev) + if err == nil { + t.Fatal("expected error when neither Event nor MessageStreamRef is set, got nil") + } + t.Logf("correctly rejected neither-set variant: %v", err) +} + +func TestAttack_SessionEventVariantNil(t *testing.T) { + ev := &TypedAgentEvent[Message]{} + err := validateAgentSessionEventIdentity(ev) + if err != nil { + t.Fatalf("nil variant should be valid, got error: %v", err) + } + t.Log("nil variant accepted correctly") +} + +func TestAttack_MessageStreamRefWrongKind(t *testing.T) { + ev := &TypedAgentEvent[Message]{ + SessionEventVariant: &SessionEventVariant[Message]{ + MessageStreamRef: &MessageStreamRef{ + EventID: "evt-1", + Kind: SessionEventSessionStatusRunning, + }, + }, + } + + err := validateAgentSessionEventIdentity(ev) + if err == nil { + t.Fatal("expected error for MessageStreamRef with non-message kind, got nil") + } + t.Logf("correctly rejected wrong kind on stream ref: %v", err) +} + +func TestAttack_ClassifySessionEventZeroValue(t *testing.T) { + ev := &SessionEvent[Message]{} + + _, err := ClassifySessionEvent(ev) + if err == nil { + t.Fatal("expected error for zero-value session event with no payload, got nil") + } + t.Logf("correctly rejected zero-value event: %v", err) +} + +func TestAttack_ClassifySessionEventNil(t *testing.T) { + _, err := ClassifySessionEvent[Message](nil) + if err == nil { + t.Fatal("expected error for nil session event, got nil") + } + t.Logf("correctly rejected nil event: %v", err) +} + +func TestAttack_ClassifySessionEventMultiplePayloads(t *testing.T) { + ev := &SessionEvent[Message]{ + Message: schema.UserMessage("hello"), + Cancel: &CancelEvent{Reason: "test"}, + } + + _, err := ClassifySessionEvent(ev) + if err == nil { + t.Fatal("expected error for event with multiple active payloads, got nil") + } + t.Logf("correctly rejected multiple-payload event: %v", err) +} + +func TestAttack_NormalizeSessionEventKindMismatch(t *testing.T) { + ev := &SessionEvent[Message]{ + Kind: SessionEventCancel, + Message: schema.UserMessage("hello"), + } + + err := NormalizeSessionEventKind(ev) + if err == nil { + t.Fatal("expected error for kind mismatch, got nil") + } + t.Logf("correctly rejected kind mismatch: %v", err) +} + +func TestAttack_NormalizeSessionEventKindUnknownKindTolerated(t *testing.T) { + unknownKinds := []SessionEventKind{ + "turn_end", + "session_started", + "custom_thing", + "future.new_kind", + } + for _, k := range unknownKinds { + ev := &SessionEvent[Message]{ + Kind: k, + } + err := NormalizeSessionEventKind(ev) + if err != nil { + t.Fatalf("unknown kind %q should be tolerated, got error: %v", k, err) + } + if ev.Kind != k { + t.Fatalf("unknown kind %q should be preserved, got %q", k, ev.Kind) + } + } +} + +func TestAttack_NormalizeSessionEventKindKnownKindMissingPayloadStillErrors(t *testing.T) { + ev := &SessionEvent[Message]{ + Kind: SessionEventMessage, + } + err := NormalizeSessionEventKind(ev) + if err == nil { + t.Fatal("expected error for known kind with missing payload, got nil") + } + t.Logf("correctly rejected known kind with missing payload: %v", err) +} + +func TestAttack_NormalizeSessionEventKindUnknownKindWithPayloadStillErrors(t *testing.T) { + ev := &SessionEvent[Message]{ + Kind: "future.new_kind", + Message: schema.UserMessage("hello"), + } + err := NormalizeSessionEventKind(ev) + if err == nil { + t.Fatal("expected error for unknown kind with recognized payload, got nil") + } + t.Logf("correctly rejected unknown kind with recognized payload: %v", err) +} + +func TestAttack_ValidateEmittedSessionEventEmptyKind(t *testing.T) { + ev := &SessionEvent[Message]{ + Message: schema.UserMessage("hello"), + } + + err := ValidateEmittedSessionEventKind(ev) + if err == nil { + t.Fatal("expected error for emitted event with empty Kind, got nil") + } + t.Logf("correctly rejected empty-kind emitted event: %v", err) +} + +func TestAttack_ValidateEmittedSessionEventNil(t *testing.T) { + err := ValidateEmittedSessionEventKind[Message](nil) + if err == nil { + t.Fatal("expected error for nil emitted event, got nil") + } + t.Logf("correctly rejected nil emitted event: %v", err) +} + +func TestAttack_SessionEventEncodeDecodeRoundtrip(t *testing.T) { + original := &SessionEvent[Message]{ + EventID: "roundtrip-1", + Timestamp: time.Date(2026, 1, 15, 10, 30, 0, 0, time.UTC), + Kind: SessionEventMessage, + Message: schema.UserMessage("roundtrip test"), + } + + data, err := encodeSessionEvent(original) + if err != nil { + t.Fatalf("encode failed: %v", err) + } + + decoded, err := decodeSessionEvent[Message](data) + if err != nil { + t.Fatalf("decode failed: %v", err) + } + + if decoded.EventID != original.EventID { + t.Errorf("EventID mismatch: got %q want %q", decoded.EventID, original.EventID) + } + if decoded.Kind != original.Kind { + t.Errorf("Kind mismatch: got %q want %q", decoded.Kind, original.Kind) + } + t.Log("encode/decode roundtrip OK") +} + +func TestAttack_SessionEventVariantPayloadEncodeDecodeRoundtrip(t *testing.T) { + original := &TypedAgentEvent[Message]{ + SessionEventVariant: &SessionEventVariant[Message]{ + SessionID: "sess-roundtrip", + Event: &SessionEvent[Message]{ + EventID: "evt-rt-1", + Timestamp: time.Date(2026, 1, 15, 10, 30, 0, 0, time.UTC), + Kind: SessionEventMessage, + Message: schema.UserMessage("variant roundtrip"), + }, + }, + } + + persistable, err := toSessionEventChecked(original) + if err != nil { + t.Fatalf("convert variant payload failed: %v", err) + } + if persistable == original.SessionEventVariant.Event { + t.Fatal("persistable event must be copied out of live SessionEventVariant") + } + + data, err := encodeSessionEvent(persistable) + if err != nil { + t.Fatalf("encode variant payload failed: %v", err) + } + + decoded, err := decodeSessionEvent[Message](data) + if err != nil { + t.Fatalf("decode variant payload failed: %v", err) + } + + if decoded.EventID != original.SessionEventVariant.Event.EventID { + t.Errorf("EventID mismatch: got %q want %q", decoded.EventID, original.SessionEventVariant.Event.EventID) + } + t.Log("variant payload encode/decode roundtrip OK") +} + +func TestAttack_AssignSessionEventIDEmptyGenerator(t *testing.T) { + emptyGen := func(_ context.Context, _ *SessionEvent[Message]) (string, error) { + return "", nil + } + + ev := &SessionEvent[Message]{ + Kind: SessionEventMessage, + Message: schema.UserMessage("test"), + } + + err := assignSessionEventID(context.Background(), ev, emptyGen) + if !errors.Is(err, ErrSessionEventIDGeneratorEmpty) { + t.Fatalf("expected ErrSessionEventIDGeneratorEmpty, got %v", err) + } + t.Logf("correctly handled empty generator: %v", err) +} + +func TestAttack_AssignSessionEventIDGeneratorError(t *testing.T) { + genErr := errors.New("generator failed") + errGen := func(_ context.Context, _ *SessionEvent[Message]) (string, error) { + return "", genErr + } + + ev := &SessionEvent[Message]{ + Kind: SessionEventMessage, + Message: schema.UserMessage("test"), + } + + err := assignSessionEventID(context.Background(), ev, errGen) + if err == nil { + t.Fatal("expected error from generator, got nil") + } + if !errors.Is(err, genErr) { + t.Fatalf("expected wrapped generator error, got %v", err) + } + t.Logf("correctly propagated generator error: %v", err) +} + +func TestAttack_ApplySessionEventMessagesDeletedEmptyIDs(t *testing.T) { + messages := []Message{schema.UserMessage("a"), schema.UserMessage("b")} + ev := &SessionEvent[Message]{ + Kind: SessionEventMessagesDeleted, + MessagesDeleted: &MessagesDeletedEvent{ + MessageIDs: []string{}, + }, + } + + err := applySessionEvent(&messages, ev) + if err == nil { + t.Fatal("expected error for empty MessageIDs, got nil") + } + t.Logf("correctly rejected empty MessageIDs: %v", err) +} + +func TestAttack_ApplySessionEventMessagesDeletedDuplicateIDs(t *testing.T) { + messages := []Message{schema.UserMessage("a")} + ev := &SessionEvent[Message]{ + Kind: SessionEventMessagesDeleted, + MessagesDeleted: &MessagesDeletedEvent{ + MessageIDs: []string{"dup", "dup"}, + }, + } + + err := applySessionEvent(&messages, ev) + if err == nil { + t.Fatal("expected error for duplicate MessageIDs, got nil") + } + t.Logf("correctly rejected duplicate MessageIDs: %v", err) +} + +func TestAttack_ApplySessionEventMessageUpdatedIdentityMismatch(t *testing.T) { + msg := schema.UserMessage("original") + EnsureMessageID(msg) + + messages := []Message{msg} + newMsg := schema.UserMessage("updated") + EnsureMessageID(newMsg) + + ev := &SessionEvent[Message]{ + Kind: SessionEventMessageUpdated, + MessageUpdated: &MessageUpdatedEvent[Message]{ + MessageID: GetMessageID(msg), + Message: newMsg, + }, + } + + err := applySessionEvent(&messages, ev) + if err == nil { + t.Log("MessageUpdated with matching ID applied OK") + } else { + t.Logf("MessageUpdated result: %v", err) + } +} + +func TestAttack_IsContextSessionEventEdgeCases(t *testing.T) { + tests := []struct { + name string + ev *SessionEvent[Message] + want bool + }{ + {"nil event", nil, false}, + {"empty event", &SessionEvent[Message]{}, false}, + {"cancel event", &SessionEvent[Message]{Cancel: &CancelEvent{}}, false}, + {"lifecycle event", &SessionEvent[Message]{Lifecycle: &LifecycleEvent{State: SessionRunStateRunning}}, false}, + {"error event", &SessionEvent[Message]{Error: &SessionErrorEvent{}}, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := isContextSessionEvent(tt.ev) + if got != tt.want { + t.Errorf("isContextSessionEvent() = %v, want %v", got, tt.want) + } + }) + } + t.Log("all isContextSessionEvent edge cases pass") +} + +func TestAttack_ToSessionEventCheckedStreamingOutput(t *testing.T) { + ev := &TypedAgentEvent[Message]{ + Output: &TypedAgentOutput[Message]{ + MessageOutput: &TypedMessageVariant[Message]{ + IsStreaming: true, + Message: nil, + }, + }, + SessionEventVariant: nil, + } + + se, err := toSessionEventChecked(ev) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if se != nil { + t.Fatalf("expected nil SessionEvent for streaming output without variant, got %v", se) + } + t.Log("streaming output without variant correctly returns nil session event") +} + +func TestAttack_StripSessionEventFields(t *testing.T) { + tests := []struct { + name string + ev *TypedAgentEvent[Message] + nil bool + }{ + {"nil event", nil, true}, + {"no variant, with output", &TypedAgentEvent[Message]{ + Output: &TypedAgentOutput[Message]{ + MessageOutput: &TypedMessageVariant[Message]{Message: schema.UserMessage("test")}, + }, + }, false}, + {"only variant", &TypedAgentEvent[Message]{ + SessionEventVariant: &SessionEventVariant[Message]{ + Event: &SessionEvent[Message]{EventID: "x", Kind: SessionEventMessage, Message: schema.UserMessage("test")}, + }, + }, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := stripSessionEventFields(tt.ev) + if tt.nil && result != nil { + t.Errorf("expected nil result, got %v", result) + } + if !tt.nil && result == nil { + t.Error("expected non-nil result, got nil") + } + if result != nil && result.SessionEventVariant != nil { + t.Error("SessionEventVariant should be stripped") + } + }) + } + t.Log("stripSessionEventFields all cases pass") +} + +func TestAttack_ClassifySpanEventBothModelAndTool(t *testing.T) { + span := &SpanEvent{ + SpanID: "span-1", + Kind: SpanKindModel, + Model: &ModelSpanMeta{}, + Tool: &ToolSpanMeta{ToolUseID: "call-1"}, + } + + _, err := classifySpanSessionEvent(span) + if err == nil { + t.Fatal("expected error when both Model and Tool are set, got nil") + } + t.Logf("correctly rejected both-model-and-tool span: %v", err) +} + +func TestAttack_ClassifySpanEventNeitherModelNorTool(t *testing.T) { + span := &SpanEvent{ + SpanID: "span-1", + Kind: SpanKindModel, + } + + _, err := classifySpanSessionEvent(span) + if err == nil { + t.Fatal("expected error when neither Model nor Tool is set, got nil") + } + t.Logf("correctly rejected no-meta span: %v", err) +} + +func TestAttack_SessionEventCancelClassification(t *testing.T) { + ev := &SessionEvent[Message]{ + Cancel: &CancelEvent{Reason: "user cancelled"}, + } + + kind, err := ClassifySessionEvent(ev) + if err != nil { + t.Fatalf("classification failed: %v", err) + } + if kind != SessionEventCancel { + t.Errorf("kind = %q, want %q", kind, SessionEventCancel) + } + t.Logf("CancelEvent classified correctly as %q", kind) +} + +func TestAttack_SessionEventInterruptClassification(t *testing.T) { + ev := &SessionEvent[Message]{ + Interrupt: &InterruptEvent{ + Contexts: []*InterruptContext{ + {InterruptID: "tool:lookup:call_1", ToolUseID: "call_1"}, + }, + }, + } + + kind, err := ClassifySessionEvent(ev) + if err != nil { + t.Fatalf("classification failed: %v", err) + } + if kind != SessionEventInterrupt { + t.Errorf("kind = %q, want %q", kind, SessionEventInterrupt) + } + t.Logf("InterruptEvent classified correctly as %q", kind) +} + +func TestAttack_SessionRollbackEventValidation(t *testing.T) { + ev := &SessionEvent[Message]{ + EventID: "rb-1", + Rollback: &SessionRollbackEvent{ + ToEventID: "target-1", + }, + } + + kind, err := ClassifySessionEvent(ev) + if err != nil { + t.Fatalf("classification failed: %v", err) + } + if kind != SessionEventRollback { + t.Errorf("kind = %q, want %q", kind, SessionEventRollback) + } + t.Logf("RollbackEvent classified correctly as %q", kind) +} + +func TestAttack_SessionRollbackEventMissingToEventID(t *testing.T) { + ev := &SessionEvent[Message]{ + EventID: "rb-1", + Rollback: &SessionRollbackEvent{}, + } + + _, err := ClassifySessionEvent(ev) + if err == nil { + t.Fatal("expected error for rollback with empty ToEventID, got nil") + } + t.Logf("correctly rejected rollback with empty ToEventID: %v", err) +} + +func TestAttack_SessionRollbackEventMissingOwnEventID(t *testing.T) { + ev := &SessionEvent[Message]{ + Rollback: &SessionRollbackEvent{ + ToEventID: "target-1", + }, + } + + _, err := ClassifySessionEvent(ev) + if err == nil { + t.Fatal("expected error for rollback with empty own EventID, got nil") + } + t.Logf("correctly rejected rollback with empty own EventID: %v", err) +} diff --git a/adk/session_timeline_test.go b/adk/session_timeline_test.go new file mode 100644 index 000000000..0338f452a --- /dev/null +++ b/adk/session_timeline_test.go @@ -0,0 +1,1567 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type sessionTimelineExtensionPayload struct { + OutcomeName string `json:"outcome_name,omitempty"` + Attempt int `json:"attempt,omitempty"` +} + +func init() { + schema.RegisterName[*sessionTimelineExtensionPayload]("_eino_adk_session_timeline_extension_payload") +} + +func requireStoredIdleStopReason(t *testing.T, raw []storedSessionEvent, want string) *SessionEvent[*schema.Message] { + t.Helper() + idleEvents := filterStoredSessionEvents(t, raw, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventSessionStatusIdle + }) + require.NotEmpty(t, idleEvents) + last := idleEvents[len(idleEvents)-1] + require.NotNil(t, last.Lifecycle) + require.NotNil(t, last.Lifecycle.StopReason) + assert.Equal(t, want, last.Lifecycle.StopReason.Type) + return last +} + +func TestSessionTimeline_ClassifyAndSerializeVariants(t *testing.T) { + now := time.Now().UTC() + spanID := uuid.NewString() + cases := []struct { + name string + se *SessionEvent[*schema.Message] + kind SessionEventKind + }{ + { + name: "lifecycle", + se: &SessionEvent[*schema.Message]{Lifecycle: &LifecycleEvent{State: SessionRunStateRunning}}, + kind: SessionEventSessionStatusRunning, + }, + { + name: "session error", + se: &SessionEvent[*schema.Message]{Error: &SessionErrorEvent{Type: SessionErrorTypeModelRetry, Message: "busy", RetryStatus: &RetryStatus{Type: "retrying"}}}, + kind: SessionEventSessionError, + }, + { + name: "span start", + se: &SessionEvent[*schema.Message]{Span: &SpanEvent{SpanID: spanID, Kind: SpanKindModel, StartedAt: now, Model: &ModelSpanMeta{}}}, + kind: SessionEventSpanModelRequestStart, + }, + { + name: "span end", + se: &SessionEvent[*schema.Message]{Span: &SpanEvent{SpanID: spanID, Kind: SpanKindModel, StartedAt: now, EndedAt: now.Add(time.Millisecond), Model: &ModelSpanMeta{}}}, + kind: SessionEventSpanModelRequestEnd, + }, + { + name: "tool span start", + se: &SessionEvent[*schema.Message]{Span: &SpanEvent{ + SpanID: spanID, Kind: SpanKindTool, StartedAt: now, + Tool: &ToolSpanMeta{ToolUseID: "call_1", Name: "lookup"}, + }}, + kind: SessionEventSpanToolCallStart, + }, + { + name: "tool span end", + se: &SessionEvent[*schema.Message]{Span: &SpanEvent{ + SpanID: spanID, Kind: SpanKindTool, StartedAt: now, EndedAt: now.Add(time.Millisecond), + Status: "ok", + Tool: &ToolSpanMeta{ToolUseID: "call_1", Name: "lookup", ToolCallStartEventID: uuid.NewString()}, + }}, + kind: SessionEventSpanToolCallEnd, + }, + { + name: "interrupt", + se: &SessionEvent[*schema.Message]{Cancel: &CancelEvent{Reason: "user"}}, + kind: SessionEventCancel, + }, + { + name: "agent interrupt", + se: &SessionEvent[*schema.Message]{Interrupt: &InterruptEvent{ + Contexts: []*InterruptContext{ + { + InterruptID: "agent:timeline-agent", + Info: "confirm?", + }, + }, + }}, + kind: SessionEventInterrupt, + }, + { + name: "extension", + se: &SessionEvent[*schema.Message]{ + Kind: SessionEventKind("x.outcome.started"), + Extension: &SessionExtensionEvent{Data: &sessionTimelineExtensionPayload{OutcomeName: "code_review"}}, + }, + kind: SessionEventKind("x.outcome.started"), + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + tc.se.EventID = uuid.NewString() + require.NoError(t, NormalizeSessionEventKind(tc.se)) + assert.Equal(t, tc.kind, tc.se.Kind) + + data, err := encodeSessionEvent(tc.se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + assert.Equal(t, tc.kind, decoded.Kind) + if tc.se.Extension != nil { + require.NotNil(t, decoded.Extension) + payload, ok := decoded.Extension.Data.(*sessionTimelineExtensionPayload) + require.True(t, ok) + assert.Equal(t, "code_review", payload.OutcomeName) + } + }) + } +} + +func TestSessionTimeline_ExtensionValidation(t *testing.T) { + t.Run("empty kind rejected", func(t *testing.T) { + err := NormalizeSessionEventKind(&SessionEvent[*schema.Message]{ + Extension: &SessionExtensionEvent{}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "must set kind") + }) + + t.Run("non extension kind rejected", func(t *testing.T) { + err := NormalizeSessionEventKind(&SessionEvent[*schema.Message]{ + Kind: SessionEventKind("outcome.started"), + Extension: &SessionExtensionEvent{}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "must start with") + }) + + t.Run("built in kind rejected", func(t *testing.T) { + err := NormalizeSessionEventKind(&SessionEvent[*schema.Message]{ + Kind: SessionEventMessage, + Extension: &SessionExtensionEvent{}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "must start with") + }) + + t.Run("built in payload cannot use extension namespace", func(t *testing.T) { + err := NormalizeSessionEventKind(&SessionEvent[*schema.Message]{ + Kind: SessionEventKind("x.message"), + Message: schema.UserMessage("hello"), + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not match payload") + }) + + t.Run("one active payload invariant", func(t *testing.T) { + err := NormalizeSessionEventKind(&SessionEvent[*schema.Message]{ + Kind: SessionEventKind("x.outcome.started"), + Message: schema.UserMessage("hello"), + Extension: &SessionExtensionEvent{}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "exactly one active payload") + }) + + t.Run("human readable typed round trip", func(t *testing.T) { + se := &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Timestamp: time.Now().UTC(), + Kind: SessionEventKind("x.outcome.grading"), + Extension: &SessionExtensionEvent{Data: &sessionTimelineExtensionPayload{Attempt: 1}}, + } + require.NoError(t, NormalizeSessionEventKind(se)) + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.Extension) + assert.Equal(t, se.Kind, decoded.Kind) + payload, ok := decoded.Extension.Data.(*sessionTimelineExtensionPayload) + require.True(t, ok) + assert.Equal(t, 1, payload.Attempt) + }) +} + +func TestSessionTimeline_ReconstructionIgnoresNonContextVariants(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "timeline-replay" + + msg := schema.UserMessage("hello") + EnsureMessageID(msg) + events := []*SessionEvent[*schema.Message]{ + {EventID: uuid.NewString(), Kind: SessionEventSessionStatusRunning, Lifecycle: &LifecycleEvent{State: SessionRunStateRunning}}, + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: msg}, + {EventID: uuid.NewString(), Kind: SessionEventSpanModelRequestStart, Span: &SpanEvent{SpanID: uuid.NewString(), Kind: SpanKindModel, StartedAt: time.Now().UTC(), Model: &ModelSpanMeta{}}}, + {EventID: uuid.NewString(), Kind: SessionEventKind("x.outcome.started"), Extension: &SessionExtensionEvent{Data: &sessionTimelineExtensionPayload{Attempt: 1}}}, + {EventID: uuid.NewString(), Kind: SessionEventInterrupt, Interrupt: &InterruptEvent{ + Contexts: []*InterruptContext{ + { + InterruptID: "agent:timeline-agent", + Info: "confirm?", + }, + }, + }}, + {EventID: uuid.NewString(), Kind: SessionEventSessionError, Error: &SessionErrorEvent{Type: "transient", RetryStatus: &RetryStatus{Type: "retrying"}}}, + {EventID: uuid.NewString(), Kind: SessionEventModelContext, ModelContext: &ModelContextEvent{}}, + } + for _, se := range events { + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 1) + assert.Equal(t, "hello", result.state.Messages[0].Content) + assert.True(t, result.state.sawModelContext) +} + +func TestSessionTimeline_AgentInterruptRoundTripPreservesContexts(t *testing.T) { + se := &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Interrupt: &InterruptEvent{ + Contexts: []*InterruptContext{ + { + InterruptID: "agent:timeline-agent;tool:lookup:call_1", + Info: "tool info", + ToolUseID: "call_1", + }, + }, + }, + } + require.NoError(t, NormalizeSessionEventKind(se)) + require.Equal(t, SessionEventInterrupt, se.Kind) + + data, err := encodeSessionEvent(se) + require.NoError(t, err) + decoded, err := decodeSessionEvent[*schema.Message](data) + require.NoError(t, err) + require.NotNil(t, decoded.Interrupt) + assert.Equal(t, SessionEventInterrupt, decoded.Kind) + require.Len(t, decoded.Interrupt.Contexts, 1) + ctx0 := decoded.Interrupt.Contexts[0] + assert.Equal(t, "agent:timeline-agent;tool:lookup:call_1", ctx0.InterruptID) + assert.Equal(t, "tool info", ctx0.Info) + assert.Equal(t, "call_1", ctx0.ToolUseID) +} + +func TestBuildInterruptEvent_ToolUseID(t *testing.T) { + contexts := []*InterruptCtx{ + { + ID: "agent:timeline-agent;tool:lookup:call_1", + Address: Address{ + {Type: AddressSegmentAgent, ID: "timeline-agent"}, + {Type: AddressSegmentTool, ID: "lookup", SubID: "call_1"}, + }, + Info: "tool info", + IsRootCause: true, + }, + } + + event := buildInterruptEvent(contexts) + require.NotNil(t, event) + require.Len(t, event.Contexts, 1) + assert.Equal(t, "agent:timeline-agent;tool:lookup:call_1", event.Contexts[0].InterruptID) + assert.Equal(t, "tool info", event.Contexts[0].Info) + assert.Equal(t, "call_1", event.Contexts[0].ToolUseID) + + // Fallback to segment ID when SubID is empty. + contexts[0].Address[1].SubID = "" + contexts[0].Address[1].ID = "legacy-call-id" + event = buildInterruptEvent(contexts) + assert.Equal(t, "legacy-call-id", event.Contexts[0].ToolUseID) +} + +func TestRunner_PersistsAgentInterruptSessionEvent(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + const checkpointID = "agent-interrupt-cp" + agent := &myAgent{ + name: "timeline-agent", + runFn: func(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(Interrupt(ctx, "confirm?")) + }() + return iter + }, + resumeFn: func(_ context.Context, _ *ResumeInfo, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + gen.Close() + return iter + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + CheckPointStore: store, + SessionID: "agent-interrupt-session", + SessionStore: store, + }) + + var liveInterruptContexts []*InterruptCtx + iter := runner.Query(ctx, "hello", WithCheckPointID(checkpointID)) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.Action != nil && event.Action.Interrupted != nil { + liveInterruptContexts = event.Action.Interrupted.InterruptContexts + } + } + require.NotEmpty(t, liveInterruptContexts) + + interrupts := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventInterrupt + }) + require.Len(t, interrupts, 1) + require.NotNil(t, interrupts[0].Interrupt) + require.Len(t, interrupts[0].Interrupt.Contexts, 1) + ctx0 := interrupts[0].Interrupt.Contexts[0] + assert.Equal(t, liveInterruptContexts[0].ID, ctx0.InterruptID) + assert.Equal(t, liveInterruptContexts[0].Info, ctx0.Info) + requireStoredIdleStopReason(t, store.events, "interrupted") + + committedIdleEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return isCommittedIdleEvent(se) + }) + assert.Empty(t, committedIdleEvents, "business interrupt should not commit") +} + +func TestSessionTimeline_ReconstructionIncludesPartialContextAfterLatestCommittedIdle(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "timeline-partial" + + committedUser := schema.UserMessage("committed user") + committedAssistant := schema.AssistantMessage("committed assistant", nil) + partialUser := schema.UserMessage("partial user") + partialAssistant := schema.AssistantMessage("partial assistant", nil) + for _, msg := range []*schema.Message{committedUser, committedAssistant, partialUser, partialAssistant} { + EnsureMessageID(msg) + } + + events := []*SessionEvent[*schema.Message]{ + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: committedUser}, + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: committedAssistant}, + {EventID: uuid.NewString(), Kind: SessionEventSessionStatusIdle, Lifecycle: &LifecycleEvent{State: SessionRunStateIdle, StopReason: &StopReason{Type: "end_turn"}}}, + {EventID: uuid.NewString(), Kind: SessionEventSessionStatusRunning, Lifecycle: &LifecycleEvent{State: SessionRunStateRunning}}, + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: partialUser}, + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: partialAssistant}, + {EventID: uuid.NewString(), Kind: SessionEventSessionError, Error: &SessionErrorEvent{Type: SessionErrorTypeModelRetry, RetryStatus: &RetryStatus{Type: "retrying"}}}, + } + for _, se := range events { + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + result, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 4) + assert.Equal(t, "committed user", result.state.Messages[0].Content) + assert.Equal(t, "committed assistant", result.state.Messages[1].Content) + assert.Equal(t, "partial user", result.state.Messages[2].Content) + assert.Equal(t, "partial assistant", result.state.Messages[3].Content) +} + +func TestSessionTimeline_ReconstructionPartialContextMissingAnchorFails(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + sid := "timeline-partial-missing-anchor" + + committedUser := schema.UserMessage("committed user") + EnsureMessageID(committedUser) + inserted := schema.SystemMessage("inserted") + EnsureMessageID(inserted) + + events := []*SessionEvent[*schema.Message]{ + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: committedUser}, + {EventID: uuid.NewString(), Kind: SessionEventSessionStatusIdle, Lifecycle: &LifecycleEvent{State: SessionRunStateIdle, StopReason: &StopReason{Type: "end_turn"}}}, + {EventID: uuid.NewString(), Kind: SessionEventMessageInserted, MessageInserted: &MessageInsertedEvent[*schema.Message]{ + Message: inserted, + BeforeMessageID: "missing-anchor", + }}, + } + for _, se := range events { + require.NoError(t, store.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + _, err := reconstructSessionState[*schema.Message](ctx, store, sid, defaultLoadPageSize) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing-anchor") +} + +func TestSessionTimeline_CommittedIdleIsReplayBoundaryOnly(t *testing.T) { + committedUser := schema.UserMessage("committed user") + partialUser := schema.UserMessage("partial user") + for _, msg := range []*schema.Message{committedUser, partialUser} { + EnsureMessageID(msg) + } + + events := []*SessionEvent[*schema.Message]{ + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: committedUser}, + {EventID: uuid.NewString(), Kind: SessionEventSessionStatusIdle, Lifecycle: &LifecycleEvent{State: SessionRunStateIdle, StopReason: &StopReason{Type: "end_turn"}}}, + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: partialUser}, + {EventID: uuid.NewString(), Kind: "turn_end"}, + } + + state, err := replayDurableContextEvents(events) + require.NoError(t, err) + require.Len(t, state.Messages, 2) + assert.Equal(t, "committed user", state.Messages[0].Content) + assert.Equal(t, "partial user", state.Messages[1].Content) +} + +func TestWithTimelineEvents_LiveExposure(t *testing.T) { + ctx := context.Background() + agent := &runnerSessionAgent{ + name: "timeline-agent", + } + + t.Run("stripped by default", func(t *testing.T) { + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: "timeline-default", SessionStore: store}) + iter := runner.Query(ctx, "hello") + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + assert.Nil(t, event.SessionEventVariant) + } + lifecycle := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventSessionStatusRunning || se.Kind == SessionEventSessionStatusIdle + }) + require.Len(t, lifecycle, 2) + requireStoredIdleStopReason(t, store.events, "end_turn") + }) + + t.Run("exposed when requested", func(t *testing.T) { + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{Agent: agent, SessionID: "timeline-visible", SessionStore: store}) + var kinds []SessionEventKind + var liveUserInput bool + iter := runner.Query(ctx, "hello", WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil { + kinds = append(kinds, event.SessionEventVariant.Event.Kind) + if event.SessionEventVariant.Event.Kind == SessionEventMessage && event.SessionEventVariant.Event.Message != nil && + event.SessionEventVariant.Event.Message.Role == schema.User && event.SessionEventVariant.Event.Message.Content == "hello" { + liveUserInput = true + } + } + } + assert.Contains(t, kinds, SessionEventSessionStatusRunning) + assert.True(t, liveUserInput, "caller input should be emitted on the live timeline") + assert.Contains(t, kinds, SessionEventSessionStatusIdle) + + storedUserInput := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventMessage && se.Message != nil && + se.Message.Role == schema.User && se.Message.Content == "hello" + }) + require.Len(t, storedUserInput, 1) + }) +} + +type extensionEventModel struct{} + +func (m *extensionEventModel) Generate(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("ok", nil), nil +} + +func (m *extensionEventModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func TestRunner_ExtensionEventSentWithTypedSendEventIsLiveAndPersisted(t *testing.T) { + ctx := context.Background() + extensionKind := SessionEventKind("x.outcome.grading") + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "extension-event-agent", + Instruction: "test", + Model: &extensionEventModel{}, + Middlewares: []AgentMiddleware{ + { + AfterChatModel: func(ctx context.Context, _ *ChatModelAgentState) error { + return SendEvent(ctx, &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: extensionKind, + Extension: &SessionExtensionEvent{ + Data: &sessionTimelineExtensionPayload{ + OutcomeName: "code_review", + Attempt: 1, + }, + }, + }, + }, + }) + }, + }, + }, + }) + require.NoError(t, err) + + t.Run("visible when timeline requested", func(t *testing.T) { + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "extension-event-session-visible", + SessionStore: store, + }) + + var liveExtension *SessionEvent[*schema.Message] + iter := runner.Query(ctx, "hello", WithTimelineEvents()) + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == extensionKind { + liveExtension = event.SessionEventVariant.Event + } + } + + require.NotNil(t, liveExtension) + require.NotEmpty(t, liveExtension.EventID) + require.NotNil(t, liveExtension.Extension) + livePayload, ok := liveExtension.Extension.Data.(*sessionTimelineExtensionPayload) + require.True(t, ok) + assert.Equal(t, "code_review", livePayload.OutcomeName) + assert.Equal(t, 1, livePayload.Attempt) + + stored := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == extensionKind + }) + require.Len(t, stored, 1) + assert.Equal(t, liveExtension.EventID, stored[0].EventID) + require.NotNil(t, stored[0].Extension) + storedPayload, ok := stored[0].Extension.Data.(*sessionTimelineExtensionPayload) + require.True(t, ok) + assert.Equal(t, "code_review", storedPayload.OutcomeName) + assert.Equal(t, 1, storedPayload.Attempt) + + var extensionIndex, idleIndex = -1, -1 + for i, payload := range store.events { + switch payload.Kind { + case extensionKind: + extensionIndex = i + case SessionEventSessionStatusIdle: + idleIndex = i + } + } + require.NotEqual(t, -1, extensionIndex) + require.NotEqual(t, -1, idleIndex) + assert.Less(t, extensionIndex, idleIndex, "extension event should enter Runner persistence before the closing idle lifecycle event") + }) + + t.Run("stripped from live stream by default", func(t *testing.T) { + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "extension-event-session-stripped", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hello") + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + assert.Nil(t, event.SessionEventVariant) + } + + stored := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == extensionKind + }) + require.Len(t, stored, 1) + }) +} + +func TestTypedSendEventOutsideExecutionIsNoop(t *testing.T) { + err := SendEvent(context.Background(), &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + Kind: SessionEventKind("x.outcome.started"), + Extension: &SessionExtensionEvent{}, + }, + }, + }) + require.NoError(t, err) +} + +func TestSessionTimeline_SpanMetaMustBeOneOf(t *testing.T) { + se := &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Span: &SpanEvent{ + SpanID: uuid.NewString(), + Kind: SpanKindModel, + StartedAt: time.Now().UTC(), + Model: &ModelSpanMeta{}, + Tool: &ToolSpanMeta{ToolUseID: "call_1"}, + }, + } + + err := NormalizeSessionEventKind(se) + require.Error(t, err) + assert.Contains(t, err.Error(), "exactly one of Model or Tool") +} + +func TestRetryTimelineEmitsRetryingError(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + generator: gen, + internalTimelineEvents: true, + }) + + emitRetryingTimeline[*schema.Message](ctx, assert.AnError) + gen.Close() + + var kinds []SessionEventKind + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + require.NotNil(t, event.SessionEventVariant.Event) + kinds = append(kinds, event.SessionEventVariant.Event.Kind) + } + + require.Equal(t, []SessionEventKind{ + SessionEventSessionError, + }, kinds) +} + +func TestModelUsageFromAssistantMapsNormalizedUsage(t *testing.T) { + usage := &schema.TokenUsage{ + PromptTokens: 10, + CompletionTokens: 5, + PromptTokenDetails: schema.PromptTokenDetails{ + CachedTokens: 7, + }, + } + msg := schema.AssistantMessage("ok", nil) + msg.ResponseMeta = &schema.ResponseMeta{Usage: usage} + + got := modelUsageFromAssistant[*schema.Message](msg) + require.NotNil(t, got) + assert.Equal(t, 10, got.InputTokens) + assert.Equal(t, 5, got.OutputTokens) + assert.Equal(t, 7, got.CacheReadInputTokens) + assert.Zero(t, got.CacheCreationInputTokens) + assert.Same(t, usage, got.Raw) +} + +func TestModelSpanEndCarriesAssistantUsage(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + generator: gen, + internalTimelineEvents: true, + }) + usage := &schema.TokenUsage{ + PromptTokens: 12, + CompletionTokens: 6, + PromptTokenDetails: schema.PromptTokenDetails{ + CachedTokens: 4, + }, + } + inner := newFakeChatModel(func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + msg := schema.AssistantMessage("ok", nil) + msg.ResponseMeta = &schema.ResponseMeta{Usage: usage, FinishReason: "stop"} + return msg, nil + }, nil) + wrapped := &typedEventSenderModel[*schema.Message]{inner: inner} + + _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + gen.Close() + + var spanEnd *SessionEvent[*schema.Message] + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == SessionEventSpanModelRequestEnd { + spanEnd = event.SessionEventVariant.Event + } + } + require.NotNil(t, spanEnd) + require.NotNil(t, spanEnd.Span.Model) + require.NotNil(t, spanEnd.Span.Model.Usage) + assert.Equal(t, 12, spanEnd.Span.Model.Usage.InputTokens) + assert.Equal(t, 6, spanEnd.Span.Model.Usage.OutputTokens) + assert.Equal(t, 4, spanEnd.Span.Model.Usage.CacheReadInputTokens) + assert.Equal(t, "stop", spanEnd.Span.Model.FinishReason) + assert.True(t, spanEnd.Span.Model.Accepted) +} + +func TestSessionTimeline_EmittedKindMustBeExplicit(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + generator: gen, + internalTimelineEvents: true, + }) + + sendSessionTimelineEvent(ctx, &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Timestamp: newEventTimestamp(), + Lifecycle: &LifecycleEvent{ + State: SessionRunStateRunning, + }, + }) + gen.Close() + + event, ok := iter.Next() + require.True(t, ok) + require.Error(t, event.Err) + assert.Contains(t, event.Err.Error(), "non-empty Kind") +} + +func TestSessionTimeline_TypedAgentEventGobRoundTripPreservesSessionEvent(t *testing.T) { + now := time.Now().UTC() + spanID := uuid.NewString() + original := &AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Timestamp: now, + Kind: SessionEventSpanToolCallStart, + Span: &SpanEvent{ + SpanID: spanID, + Kind: SpanKindTool, + Name: "tool_call", + StartedAt: now, + Tool: &ToolSpanMeta{ToolUseID: "call_1", Name: "lookup"}, + }, + }, + }, + } + + var buf bytes.Buffer + require.NoError(t, gob.NewEncoder(&buf).Encode(original)) + + var decoded AgentEvent + require.NoError(t, gob.NewDecoder(&buf).Decode(&decoded)) + require.NotNil(t, decoded.SessionEventVariant.Event) + assert.Equal(t, original.SessionEventVariant.Event.EventID, decoded.SessionEventVariant.Event.EventID) + assert.Equal(t, SessionEventSpanToolCallStart, decoded.SessionEventVariant.Event.Kind) +} + +func TestModelSpanMetaFromContextPopulatesFailoverAndModelFields(t *testing.T) { + parentSpanID := uuid.NewString() + ctx := context.Background() + ctx = typedSetFailoverCurrentModel[*schema.Message](ctx, newFakeChatModel(nil, nil)) + ctx = withFailoverTimeline(ctx, parentSpanID, 3) + + started := newEventTimestamp() + start := newModelSpanStartEvent[*schema.Message](ctx, uuid.NewString(), started, model.WithModel("claude-sonnet")) + require.NotNil(t, start.Span) + require.NotNil(t, start.Span.Model) + assert.Equal(t, parentSpanID, start.Span.ParentSpanID) + assert.Equal(t, "fake_chat_model", start.Span.Model.Provider) + assert.Equal(t, "claude-sonnet", start.Span.Model.Model) + assert.Equal(t, 3, start.Span.Model.Attempt) + + end := newModelSpanEndEvent[*schema.Message]( + ctx, + modelSpanEndEventInput[*schema.Message]{ + spanID: start.Span.SpanID, + startEventID: start.EventID, + started: started, + ended: started.Add(time.Millisecond), + msg: schema.AssistantMessage("ok", nil), + accepted: true, + }, + model.WithModel("claude-sonnet"), + ) + require.NotNil(t, end.Span) + require.NotNil(t, end.Span.Model) + assert.Equal(t, parentSpanID, end.Span.ParentSpanID) + assert.Equal(t, start.EventID, end.Span.Model.ModelRequestStartEventID) + assert.Equal(t, 3, end.Span.Model.Attempt) +} + +func TestRetryTimelineUsesRejectReasonMessage(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + generator: gen, + internalTimelineEvents: true, + }) + + var calls int + inner := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + calls++ + if calls == 1 { + return schema.AssistantMessage("bad", nil), nil + } + return schema.AssistantMessage("ok", nil), nil + }, nil) + wrapper := newTypedRetryModelWrapper[*schema.Message](inner, &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad" { + return &RetryDecision{Retry: true, RejectReason: "policy rejected"} + } + return &RetryDecision{} + }, + BackoffFunc: func(context.Context, int) time.Duration { return 0 }, + }) + + msg, err := wrapper.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + assert.Equal(t, "ok", msg.Content) + gen.Close() + + var found bool + for { + event, ok := iter.Next() + if !ok { + break + } + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == SessionEventSessionError { + require.NotNil(t, event.SessionEventVariant.Event.Error) + assert.Equal(t, "policy rejected", event.SessionEventVariant.Event.Error.Message) + found = true + } + } + assert.True(t, found) +} + +func TestFailoverTimelineLinksAttemptsAndEmitsSessionErrors(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + modelErr := errors.New("first failed") + m1 := newFakeChatModel(func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return nil, modelErr + }, nil) + m2 := newFakeChatModel(func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("ok", nil), nil + }, nil) + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + GetFailoverModel: func(context.Context, *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + }, + }) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + generator: gen, + internalTimelineEvents: true, + failoverLastSuccessModel: m1, + }) + + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}, model.WithModel("logical-model")) + require.NoError(t, err) + assert.Equal(t, "ok", msg.Content) + gen.Close() + + var starts []*SessionEvent[*schema.Message] + var failoverErrors []*SessionEvent[*schema.Message] + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant == nil || event.SessionEventVariant.Event == nil { + continue + } + switch event.SessionEventVariant.Event.Kind { + case SessionEventSpanModelRequestStart: + starts = append(starts, event.SessionEventVariant.Event) + case SessionEventSessionError: + if event.SessionEventVariant.Event.Error != nil && event.SessionEventVariant.Event.Error.Type == SessionErrorTypeModelFailover { + failoverErrors = append(failoverErrors, event.SessionEventVariant.Event) + } + } + } + require.Len(t, starts, 2) + require.NotEmpty(t, starts[0].Span.ParentSpanID) + assert.Equal(t, starts[0].Span.ParentSpanID, starts[1].Span.ParentSpanID) + assert.Equal(t, 1, starts[0].Span.Model.Attempt) + assert.Equal(t, 2, starts[1].Span.Model.Attempt) + assert.Equal(t, "logical-model", starts[0].Span.Model.Model) + require.Len(t, failoverErrors, 1) + assert.Equal(t, "retrying", failoverErrors[0].Error.RetryStatus.Type) +} + +func TestSessionTimeline_EventIDMismatchRejectedAtPersistenceBoundary(t *testing.T) { + now := time.Now().UTC() + err := validateAgentSessionEventIdentity(&AgentEvent{ + SessionEventVariant: &SessionEventVariant[*schema.Message]{ + Event: &SessionEvent[*schema.Message]{ + EventID: uuid.NewString(), + Timestamp: now, + Kind: SessionEventSpanToolCallStart, + Span: &SpanEvent{ + SpanID: uuid.NewString(), + Kind: SpanKindTool, + StartedAt: now, + Tool: &ToolSpanMeta{ToolUseID: "call_1"}, + }, + }, + MessageStreamRef: &MessageStreamRef{ + EventID: uuid.NewString(), + Timestamp: now, + Kind: SessionEventMessage, + }, + }, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "exactly one") +} + +func TestSessionTimeline_NormalizeAgentSessionEventMaterializesEnvelope(t *testing.T) { + now := time.Now().UTC() + makeToolStartSpan := func() *SessionEvent[*schema.Message] { + return &SessionEvent[*schema.Message]{ + Kind: SessionEventSpanToolCallStart, + Span: &SpanEvent{ + SpanID: uuid.NewString(), + Kind: SpanKindTool, + StartedAt: now, + Tool: &ToolSpanMeta{ToolUseID: "call_1"}, + }, + } + } + + t.Run("id empty", func(t *testing.T) { + original := makeToolStartSpan() + event := &AgentEvent{SessionEventVariant: &SessionEventVariant[*schema.Message]{Event: original}} + se, err := normalizeAgentSessionEvent(event) + require.NoError(t, err) + require.NotEmpty(t, se.EventID) + assert.Equal(t, se.EventID, event.SessionEventVariant.Event.EventID) + require.False(t, se.Timestamp.IsZero()) + assert.Equal(t, se.Timestamp, event.SessionEventVariant.Event.Timestamp) + assert.Empty(t, original.EventID) + }) + + t.Run("model context event normalizes without mutation", func(t *testing.T) { + original := &SessionEvent[*schema.Message]{Kind: SessionEventModelContext, ModelContext: &ModelContextEvent{}} + event := &AgentEvent{SessionEventVariant: &SessionEventVariant[*schema.Message]{Event: original}} + se, err := normalizeAgentSessionEvent(event) + require.NoError(t, err) + require.NotNil(t, se.ModelContext) + require.NotNil(t, event.SessionEventVariant.Event.ModelContext) + }) +} + +func TestRetryOnlyModelSpansHaveNoParentSpanID(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + var calls int + modelErr := errors.New("retry me") + inner := newFakeChatModel(func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + calls++ + if calls == 1 { + return nil, modelErr + } + return schema.AssistantMessage("ok", nil), nil + }, nil) + wrapped := buildModelWrappers[*schema.Message](inner, &modelWrapperConfig{ + retryConfig: &ModelRetryConfig{ + MaxRetries: 1, + BackoffFunc: func(context.Context, int) time.Duration { + return 0 + }, + }, + }) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + generator: gen, + internalTimelineEvents: true, + }) + + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + assert.Equal(t, "ok", msg.Content) + gen.Close() + + var starts []*SessionEvent[*schema.Message] + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == SessionEventSpanModelRequestStart { + starts = append(starts, event.SessionEventVariant.Event) + } + } + require.NotEmpty(t, starts) + for _, start := range starts { + assert.Empty(t, start.Span.ParentSpanID) + } +} + +func TestRetryAndFailoverTimelineKeepsDistinctErrorTypes(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + m1 := newFakeChatModel(func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return nil, errors.New("primary failed") + }, nil) + m2 := newFakeChatModel(func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("fallback ok", nil), nil + }, nil) + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: &ModelRetryConfig{ + MaxRetries: 1, + BackoffFunc: func(context.Context, int) time.Duration { + return 0 + }, + }, + failoverConfig: &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + GetFailoverModel: func(context.Context, *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + }, + }) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + generator: gen, + internalTimelineEvents: true, + failoverLastSuccessModel: m1, + }) + + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + assert.Equal(t, "fallback ok", msg.Content) + gen.Close() + + var retryExhausted bool + var failoverRetrying bool + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant == nil || event.SessionEventVariant.Event == nil || event.SessionEventVariant.Event.Kind != SessionEventSessionError || event.SessionEventVariant.Event.Error == nil { + continue + } + switch event.SessionEventVariant.Event.Error.Type { + case SessionErrorTypeModelRetry: + if event.SessionEventVariant.Event.Error.RetryStatus != nil && event.SessionEventVariant.Event.Error.RetryStatus.Type == "exhausted" { + retryExhausted = true + } + case SessionErrorTypeModelFailover: + if event.SessionEventVariant.Event.Error.RetryStatus != nil && event.SessionEventVariant.Event.Error.RetryStatus.Type == "retrying" { + failoverRetrying = true + } + } + } + assert.True(t, retryExhausted) + assert.True(t, failoverRetrying) +} + +type timelineErrorAgent struct { + name string + err error +} + +func (a *timelineErrorAgent) Name(context.Context) string { + return a.name +} + +func (a *timelineErrorAgent) Description(context.Context) string { + return "timeline error agent" +} + +func (a *timelineErrorAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(&AgentEvent{Err: a.err}) + }() + return iter +} + +func TestRunnerTimelineRetryExhaustedStopReason(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: &timelineErrorAgent{name: "retry-exhausted", err: &RetryExhaustedError{LastErr: errors.New("still failing"), TotalRetries: 1}}, + SessionID: "timeline-retry-exhausted", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hi") + for { + if _, ok := iter.Next(); !ok { + break + } + } + + requireStoredIdleStopReason(t, store.events, "retries_exhausted") + + turnEnds := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return isCommittedIdleEvent(se) + }) + assert.Empty(t, turnEnds, "retry exhaustion should not commit") +} + +func TestRunnerTimelineFailedStopReason(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: &timelineErrorAgent{name: "failed", err: errors.New("boom")}, + SessionID: "timeline-failed", + SessionStore: store, + }) + + iter := runner.Query(ctx, "hi") + var gotErrs []error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + gotErrs = append(gotErrs, event.Err) + } + } + require.NotEmpty(t, gotErrs) + assert.EqualError(t, gotErrs[0], "boom") + for _, err := range gotErrs { + assert.NotContains(t, err.Error(), "missing committed idle") + } + + sessionErrors := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventSessionError + }) + require.NotEmpty(t, sessionErrors) + require.NotNil(t, sessionErrors[len(sessionErrors)-1].Error) + assert.Equal(t, SessionErrorTypeFatal, sessionErrors[len(sessionErrors)-1].Error.Type) + assert.Equal(t, "boom", sessionErrors[len(sessionErrors)-1].Error.Message) + requireStoredIdleStopReason(t, store.events, "failed") + + committedIdleEvents := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return isCommittedIdleEvent(se) + }) + assert.Empty(t, committedIdleEvents, "failed turn should not commit") +} + +func TestRunnerTimelineModelCallFatalDoesNotRequireCommitMarker(t *testing.T) { + ctx := context.Background() + modelErr := errors.New("model exploded") + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "fatal-model-agent", + Model: newFakeChatModel(func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return nil, modelErr + }, nil), + }) + require.NoError(t, err) + + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "timeline-fatal-model", + SessionStore: store, + }) + + var gotErrs []error + iter := runner.Query(ctx, "hi") + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + gotErrs = append(gotErrs, event.Err) + } + } + + require.NotEmpty(t, gotErrs) + assert.ErrorIs(t, gotErrs[0], modelErr) + for _, err := range gotErrs { + assert.NotContains(t, err.Error(), "missing committed idle") + } + + sessionErrors := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventSessionError + }) + require.NotEmpty(t, sessionErrors) + require.NotNil(t, sessionErrors[len(sessionErrors)-1].Error) + assert.Equal(t, SessionErrorTypeFatal, sessionErrors[len(sessionErrors)-1].Error.Type) + assert.Contains(t, sessionErrors[len(sessionErrors)-1].Error.Message, modelErr.Error()) + requireStoredIdleStopReason(t, store.events, "failed") + + turnEnds := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return isCommittedIdleEvent(se) + }) + assert.Empty(t, turnEnds, "fatal model call should not commit") +} + +func TestRunnerTimelineCancelStopReasonAndUserInterruptPersisted(t *testing.T) { + ctx := context.Background() + store := newSessionHelperStore() + started := make(chan struct{}) + release := make(chan struct{}) + + agent := &myAgent{ + name: "timeline-cancel", + runFn: func(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + close(started) + <-release + gen.Send(Interrupt(ctx, "cancel point")) + }() + return iter + }, + } + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + CheckPointStore: store, + SessionID: "timeline-cancel", + SessionStore: store, + }) + cancelOpt, cancelFn := WithCancel() + iter := runner.Query(ctx, "hi", cancelOpt, WithCheckPointID("timeline-cancel-cp")) + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("agent did not start") + } + cancelHandle, contributed := cancelFn(WithAgentCancelMode(CancelImmediate)) + require.True(t, contributed) + close(release) + + var sawCancelErr bool + for { + event, ok := iter.Next() + if !ok { + break + } + var cancelErr *CancelError + if event.Err != nil && errors.As(event.Err, &cancelErr) { + sawCancelErr = true + } + } + require.True(t, sawCancelErr, "expected CancelError in event stream") + require.NoError(t, cancelHandle.Wait()) + + userInterrupts := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventCancel + }) + require.Len(t, userInterrupts, 1) + assert.Equal(t, SessionEventKind("cancel"), userInterrupts[0].Kind) + require.NotNil(t, userInterrupts[0].Cancel) + assert.Equal(t, "cancelled", userInterrupts[0].Cancel.Reason) + requireStoredIdleStopReason(t, store.events, "cancelled") + + turnEnds := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return isCommittedIdleEvent(se) + }) + assert.Empty(t, turnEnds, "cancelled turn should not commit") +} + +func TestToolSpan_PersistedAroundToolCallAndLinksToMessages(t *testing.T) { + ctx := context.Background() + testTool := &invokableTestTool{name: "tool_span_tool", result: "tool result"} + mockModel := &mockToolCallingModel{toolCallName: "tool_span_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ToolSpanAgent", + Description: "tool span agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + }) + require.NoError(t, err) + + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "tool-span-around", + SessionStore: store, + }) + iter := runner.Query(ctx, "go", WithTimelineEvents()) + var liveToolEnd *SessionEvent[*schema.Message] + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + if event.SessionEventVariant != nil && event.SessionEventVariant.Event != nil && event.SessionEventVariant.Event.Kind == SessionEventSpanToolCallEnd { + liveToolEnd = event.SessionEventVariant.Event + } + } + require.NotNil(t, liveToolEnd, "expected live tool_call_end span emission") + require.NotNil(t, liveToolEnd.Span) + require.NotNil(t, liveToolEnd.Span.Tool) + assert.Equal(t, "tool_span_tool", liveToolEnd.Span.Tool.Name) + assert.Equal(t, "ok", liveToolEnd.Span.Status) + + stored := filterStoredSessionEvents(t, store.events, func(_ *SessionEvent[*schema.Message]) bool { return true }) + var ( + assistantMsgEventID string + toolResultEventID string + toolStart *SessionEvent[*schema.Message] + toolEnd *SessionEvent[*schema.Message] + toolUseObservations int + ) + for _, se := range stored { + switch se.Kind { + case SessionEventMessage: + if se.Message != nil && se.Message.Role == schema.Assistant && len(se.Message.ToolCalls) > 0 { + assistantMsgEventID = se.EventID + } + if se.Message != nil && se.Message.Role == schema.Tool { + toolResultEventID = se.EventID + } + case SessionEventSpanToolCallStart: + toolStart = se + case SessionEventSpanToolCallEnd: + toolEnd = se + case "agent.tool_use", "agent.tool_result", "agent.thinking": + toolUseObservations++ + } + } + + require.NotNil(t, toolStart, "expected tool_call_start span") + require.NotNil(t, toolEnd, "expected tool_call_end span") + require.NotNil(t, toolStart.Span.Tool) + require.NotNil(t, toolEnd.Span.Tool) + assert.Equal(t, "tool_span_tool", toolStart.Span.Tool.Name) + assert.Equal(t, "tool_span_tool", toolEnd.Span.Tool.Name) + assert.Equal(t, "tc-1", toolStart.Span.Tool.ToolUseID) + assert.Equal(t, toolStart.EventID, toolEnd.Span.Tool.ToolCallStartEventID) + assert.Equal(t, "ok", toolEnd.Span.Status) + assert.Equal(t, 0, toolUseObservations, "no observation kinds should be persisted") + + if assistantMsgEventID != "" { + assert.Equal(t, assistantMsgEventID, toolStart.Span.Tool.AssistantMessageEventID) + assert.Equal(t, assistantMsgEventID, toolEnd.Span.Tool.AssistantMessageEventID) + } + if toolResultEventID != "" { + assert.Equal(t, toolResultEventID, toolEnd.Span.Tool.ToolResultMessageEventID) + } + assert.NotEmpty(t, toolStart.Span.ParentSpanID, "parent should be the model request span") + assert.Equal(t, toolStart.Span.ParentSpanID, toolEnd.Span.ParentSpanID) +} + +// TestSessionEventIDGenerator_CustomToolResultBusinessID 验证：configured +// generator 看到 tool result message 草稿时返回业务 ID，持久化的 message +// EventID 与对应 tool span end 的 ToolResultMessageEventID 必须等于该业务 ID +// （§8 CustomToolResult 验收）。 +func TestSessionEventIDGenerator_CustomToolResultBusinessID(t *testing.T) { + ctx := context.Background() + testTool := &invokableTestTool{name: "tool_span_tool", result: "tool result"} + mockModel := &mockToolCallingModel{toolCallName: "tool_span_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ToolSpanGenAgent", + Description: "tool span agent with id generator", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + }) + require.NoError(t, err) + + const toolResultBusinessID = "custom-result-id" + gen := func(_ context.Context, e *SessionEvent[*schema.Message]) (string, error) { + if e != nil && e.Kind == SessionEventMessage && e.Message != nil && e.Message.Role == schema.Tool { + return toolResultBusinessID, nil + } + return DefaultSessionEventIDGenerator[*schema.Message](ctx, e) + } + + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "tool-result-business-id", + SessionStore: store, + SessionConfig: &SessionConfig[*schema.Message]{ + EventIDGenerator: gen, + }, + }) + iter := runner.Query(ctx, "go") + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + } + + stored := filterStoredSessionEvents(t, store.events, func(_ *SessionEvent[*schema.Message]) bool { return true }) + var ( + toolResultMsg *SessionEvent[*schema.Message] + toolEnd *SessionEvent[*schema.Message] + ) + for _, se := range stored { + switch { + case se.Kind == SessionEventMessage && se.Message != nil && se.Message.Role == schema.Tool: + toolResultMsg = se + case se.Kind == SessionEventSpanToolCallEnd: + toolEnd = se + } + } + require.NotNil(t, toolResultMsg, "expected persisted tool result message") + require.NotNil(t, toolEnd, "expected tool_call_end span") + require.NotNil(t, toolEnd.Span) + require.NotNil(t, toolEnd.Span.Tool) + + assert.Equal(t, toolResultBusinessID, toolResultMsg.EventID, + "tool result message must adopt the generator-supplied business ID") + assert.Equal(t, toolResultBusinessID, toolEnd.Span.Tool.ToolResultMessageEventID, + "tool span end ToolResultMessageEventID must match the tool result message business ID") +} + +type kindsRecordingStore struct { + inner *sessionHelperStore + recordedKinds [][]SessionEventKind +} + +func (s *kindsRecordingStore) loadEvents(ctx context.Context, opts *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + if opts != nil { + s.recordedKinds = append(s.recordedKinds, opts.Kinds) + } + return s.inner.LoadEventsForSession(ctx, "", opts) +} + +func (s *kindsRecordingStore) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return s.inner.AppendEventsForSession(ctx, "", events) +} + +func (s *kindsRecordingStore) close(context.Context) error { return nil } + +func TestSessionTimeline_ReconstructionUsesKindFilter(t *testing.T) { + ctx := context.Background() + inner := newSessionHelperStore() + wrapper := &kindsRecordingStore{inner: inner} + sid := "timeline-kind-filter" + + msg1 := schema.UserMessage("hello") + EnsureMessageID(msg1) + msg2 := schema.AssistantMessage("world", nil) + EnsureMessageID(msg2) + + events := []*SessionEvent[*schema.Message]{ + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: msg1}, + {EventID: uuid.NewString(), Kind: SessionEventMessage, Message: msg2}, + {EventID: uuid.NewString(), Kind: SessionEventSpanModelRequestStart, Span: &SpanEvent{SpanID: uuid.NewString(), Kind: SpanKindModel, StartedAt: time.Now().UTC(), Model: &ModelSpanMeta{}}}, + {EventID: uuid.NewString(), Kind: SessionEventModelContext, ModelContext: &ModelContextEvent{}}, + } + for _, se := range events { + require.NoError(t, inner.AppendEventsForSession(ctx, sid, []*SessionEvent[*schema.Message]{se})) + } + + result, err := reconstructSessionState[*schema.Message](ctx, wrapper, sid, defaultLoadPageSize) + require.NoError(t, err) + + // All recorded Kinds slices should equal sessionReplayEventKinds. + require.NotEmpty(t, wrapper.recordedKinds) + for _, kinds := range wrapper.recordedKinds { + assert.Equal(t, sessionReplayEventKinds, kinds) + } + + // Verify reconstruction result. + require.NotNil(t, result) + require.NotNil(t, result.state) + require.Len(t, result.state.Messages, 2) + assert.Equal(t, "hello", result.state.Messages[0].Content) + assert.Equal(t, "world", result.state.Messages[1].Content) + assert.True(t, result.state.sawModelContext) +} + +func TestToolSpan_StreamableToolEmitsEndAfterEOF(t *testing.T) { + ctx := context.Background() + streamTool := &streamableTestTool{name: "stream_span_tool", result: "stream chunk"} + mockModel := &mockToolCallingModel{toolCallName: "stream_span_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamSpanAgent", + Description: "stream span agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{streamTool}}, + }, + }) + require.NoError(t, err) + + store := newSessionHelperStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + SessionID: "tool-span-stream", + SessionStore: store, + }) + iter := runner.Query(ctx, "stream go") + for { + event, ok := iter.Next() + if !ok { + break + } + require.NoError(t, event.Err) + } + + stored := filterStoredSessionEvents(t, store.events, func(se *SessionEvent[*schema.Message]) bool { + return se.Kind == SessionEventSpanToolCallStart || se.Kind == SessionEventSpanToolCallEnd + }) + require.Len(t, stored, 2) + assert.Equal(t, SessionEventSpanToolCallStart, stored[0].Kind) + assert.Equal(t, SessionEventSpanToolCallEnd, stored[1].Kind) + assert.Equal(t, "ok", stored[1].Span.Status) + require.NotNil(t, stored[1].Span.Tool) + assert.NotEmpty(t, stored[1].Span.Tool.ToolResultMessageEventID) +} diff --git a/adk/turn_loop.go b/adk/turn_loop.go index c2fe185b4..520bc8533 100644 --- a/adk/turn_loop.go +++ b/adk/turn_loop.go @@ -38,6 +38,22 @@ const ( stopCommitted ) +// TurnLoopResumeDecision is returned by GenResume to choose what to do with a +// pending runner checkpoint. +type TurnLoopResumeDecision int + +const ( + // TurnLoopResumeDecisionResume resumes the suspended runner checkpoint. + TurnLoopResumeDecisionResume TurnLoopResumeDecision = iota + // TurnLoopResumeDecisionStartNewTurn abandons the checkpoint and starts a + // fresh Runner.Run turn using GenResumeResult.Input. + TurnLoopResumeDecisionStartNewTurn +) + +var ( + ErrTurnLoopStopped = errors.New("adk: turn loop stopped") +) + type preemptTurnPhase uint8 const ( @@ -566,21 +582,18 @@ type TurnLoopConfig[T any, M MessageType] struct { // Required. GenInput func(ctx context.Context, loop *TurnLoop[T, M], items []T) (*GenInputResult[T, M], error) - // GenResume is called at most once during Run(). When CheckpointID is - // configured, Run() queries Store for the checkpoint: - // - If the checkpoint contains runner state (i.e. an agent was interrupted - // or canceled mid-turn), Run() calls GenResume to plan a resume turn. - // - Otherwise (no checkpoint, or between-turns checkpoint), GenResume is - // never called and the loop proceeds via GenInput. + // GenResume is called when the loop restores a pending runner checkpoint + // from Store and needs user policy to continue. // // It receives: // - interruptedItems: the items being processed when the prior run was interrupted / canceled - // - unhandledItems: items buffered but not processed when the prior run exited - // - newItems: items that were Push()-ed before Run() was called + // - unhandledItems: normal items buffered but not processed + // - newItems: items treated as resume intent (pre-run Push items for legacy + // checkpoints with no persisted ResumeItems, or persisted ResumeItems for + // checkpoints that already had accepted resume items). // - // It returns a GenResumeResult describing how to resume the interrupted agent - // turn (optional ResumeParams) and how to manipulate the buffer - // (Consumed/Remaining) before continuing. + // It returns a GenResumeResult choosing whether to resume the suspended + // runner checkpoint or abandon it and start a fresh turn. GenResume func(ctx context.Context, loop *TurnLoop[T, M], interruptedItems, unhandledItems, newItems []T) (*GenResumeResult[T, M], error) // PrepareAgent returns an Agent configured to handle the consumed items. @@ -630,6 +643,13 @@ type TurnLoopConfig[T any, M MessageType] struct { // same CheckpointID. On clean exit (no checkpoint saved), the existing // checkpoint under CheckpointID is deleted to prevent stale resumption. CheckpointID string + + // Session fields are passed through to the internal Runner used by TurnLoop. + // They let the Runner reconstruct session context for both fresh turns and + // resumed turns from a persisted checkpoint. + SessionID string + SessionStore SessionEventStore[M] + SessionConfig *SessionConfig[M] } // GenInputResult contains the result of GenInput processing. @@ -680,6 +700,13 @@ type GenResumeResult[T any, M MessageType] struct { // ResumeParams are optional parameters for resuming an interrupted agent. ResumeParams *ResumeParams + // Decision selects whether to resume the suspended checkpoint or abandon it + // and start a fresh turn. The zero value resumes for compatibility. + Decision TurnLoopResumeDecision + + // Input is required when Decision is TurnLoopResumeDecisionStartNewTurn. + Input *TypedAgentInput[M] + // Consumed are the items selected for this resumed turn. // They are removed from the buffer and passed to PrepareAgent. Consumed []T @@ -687,19 +714,20 @@ type GenResumeResult[T any, M MessageType] struct { // Remaining are the items to keep in the buffer for a future turn. // TurnLoop pushes Remaining back into the buffer before resuming the agent. // - // Items from (interruptedItems, unhandledItems, newItems) that are in neither Consumed + // Items from (interruptedItems, unhandledItems, resume items) that are in neither Consumed // nor Remaining are dropped by the loop. Remaining []T } type turnRunSpec[T any, M MessageType] struct { - runCtx context.Context - input *TypedAgentInput[M] - runOpts []AgentRunOption - resumeParams *ResumeParams - isResume bool - consumed []T - resumeBytes []byte + runCtx context.Context + input *TypedAgentInput[M] + runOpts []AgentRunOption + resumeParams *ResumeParams + isResume bool + consumed []T + resumeCheckpointID string + resumeBytes []byte } type turnPlan[T any, M MessageType] struct { @@ -746,7 +774,7 @@ func (l *TurnLoop[T, M]) planTurn( if l.config.GenResume == nil { return nil, errors.New("GenResume is required for resume") } - resumeResult, err := l.config.GenResume(ctx, l, pr.interrupted, pr.unhandled, pr.newItems) + resumeResult, err := l.config.GenResume(ctx, l, pr.interrupted, pr.unhandled, pr.resumeItems) if err != nil { return nil, err } @@ -757,18 +785,48 @@ func (l *TurnLoop[T, M]) planTurn( if resumeResult.RunCtx != nil { turnCtx = resumeResult.RunCtx } - return &turnPlan[T, M]{ - turnCtx: turnCtx, - remaining: resumeResult.Remaining, - spec: &turnRunSpec[T, M]{ - runCtx: resumeResult.RunCtx, - runOpts: resumeResult.RunOpts, - resumeParams: resumeResult.ResumeParams, - isResume: true, - consumed: resumeResult.Consumed, - resumeBytes: pr.resumeBytes, - }, - }, nil + switch resumeResult.Decision { + case TurnLoopResumeDecisionResume: + if resumeResult.Input != nil { + return nil, errors.New("GenResumeResult.Input must be nil when resuming") + } + if len(pr.resumeBytes) == 0 { + return nil, errors.New("resume checkpoint is empty") + } + resumeCheckpointID := pr.resumeCheckpointID + if resumeCheckpointID == "" { + resumeCheckpointID = bridgeCheckpointID + } + return &turnPlan[T, M]{ + turnCtx: turnCtx, + remaining: resumeResult.Remaining, + spec: &turnRunSpec[T, M]{ + runCtx: resumeResult.RunCtx, + runOpts: resumeResult.RunOpts, + resumeParams: resumeResult.ResumeParams, + isResume: true, + consumed: resumeResult.Consumed, + resumeCheckpointID: resumeCheckpointID, + resumeBytes: pr.resumeBytes, + }, + }, nil + case TurnLoopResumeDecisionStartNewTurn: + if resumeResult.Input == nil { + return nil, errors.New("GenResumeResult.Input is nil for fresh turn") + } + return &turnPlan[T, M]{ + turnCtx: turnCtx, + remaining: resumeResult.Remaining, + spec: &turnRunSpec[T, M]{ + runCtx: resumeResult.RunCtx, + input: resumeResult.Input, + runOpts: resumeResult.RunOpts, + consumed: resumeResult.Consumed, + }, + }, nil + default: + return nil, fmt.Errorf("unknown GenResume decision: %d", resumeResult.Decision) + } } // InterruptError is the ExitReason when the TurnLoop exits due to a business @@ -916,10 +974,12 @@ type TurnLoop[T any, M MessageType] struct { interruptedItems []T checkPointRunnerBytes []byte + checkPointRunnerID string interruptContexts []*InterruptCtx capturedCancelErr *CancelError pendingResume *turnLoopPendingResume[T] + resumeMu sync.Mutex loadCheckpointID string @@ -940,12 +1000,14 @@ func (l *TurnLoop[T, M]) appendLate(item T) { } type turnLoopCheckpoint[T any] struct { - RunnerCheckpoint []byte + RunnerCheckpointID string + RunnerCheckpoint []byte // HasRunnerState reports whether RunnerCheckpoint contains resumable runner state. // It is false for "between turns" checkpoints where no agent execution was // interrupted (e.g. Stop() before the first turn or between turns). HasRunnerState bool UnhandledItems []T + ResumeItems []T CanceledItems []T // gob-compat: kept as CanceledItems for deserialization of existing checkpoints } @@ -986,6 +1048,18 @@ func (l *TurnLoop[T, M]) deleteTurnLoopCheckpoint(ctx context.Context, checkPoin return nil } +func (l *TurnLoop[T, M]) deleteLoadedCheckpointAfterSuccessfulResume(ctx context.Context, runErr error, isResume bool, hasInterrupt bool) error { + if runErr != nil || !isResume || hasInterrupt || l.loadCheckpointID == "" { + return runErr + } + checkpointID := l.loadCheckpointID + if err := l.deleteTurnLoopCheckpoint(ctx, checkpointID); err != nil { + return fmt.Errorf("failed to delete consumed checkpoint[%s] after resume: %w", checkpointID, err) + } + l.loadCheckpointID = "" + return nil +} + func (l *TurnLoop[T, M]) tryLoadCheckpoint(ctx context.Context) error { checkPointID := l.config.CheckpointID if checkPointID == "" || l.config.Store == nil { @@ -1018,11 +1092,27 @@ func (l *TurnLoop[T, M]) tryLoadCheckpoint(ctx context.Context) error { l.buffer.PushFront(newItems) return fmt.Errorf("checkpoint[%s] has runner state but bytes are empty", checkPointID) } + resumeCheckpointID := cp.RunnerCheckpointID + if resumeCheckpointID == "" { + resumeCheckpointID = bridgeCheckpointID + } + resumeItems := append([]T{}, cp.ResumeItems...) + resumeSubmitted := len(resumeItems) > 0 + if !resumeSubmitted { + resumeItems = append(resumeItems, newItems...) + } else { + unhandled := make([]T, 0, len(cp.UnhandledItems)+len(newItems)) + unhandled = append(unhandled, cp.UnhandledItems...) + unhandled = append(unhandled, newItems...) + cp.UnhandledItems = unhandled + } l.pendingResume = &turnLoopPendingResume[T]{ - interrupted: append([]T{}, cp.CanceledItems...), - unhandled: append([]T{}, cp.UnhandledItems...), - newItems: append([]T{}, newItems...), - resumeBytes: append([]byte{}, cp.RunnerCheckpoint...), + interrupted: append([]T{}, cp.CanceledItems...), + unhandled: append([]T{}, cp.UnhandledItems...), + resumeItems: resumeItems, + resumeSubmitted: resumeSubmitted, + resumeCheckpointID: resumeCheckpointID, + resumeBytes: append([]byte{}, cp.RunnerCheckpoint...), } } else { items := make([]T, 0, len(cp.UnhandledItems)+len(newItems)) @@ -1035,10 +1125,12 @@ func (l *TurnLoop[T, M]) tryLoadCheckpoint(ctx context.Context) error { } type turnLoopPendingResume[T any] struct { - interrupted []T - unhandled []T - newItems []T - resumeBytes []byte + interrupted []T + unhandled []T + resumeItems []T + resumeSubmitted bool + resumeCheckpointID string + resumeBytes []byte } // SafePoint describes at which boundary the agent may be cancelled. @@ -1567,6 +1659,131 @@ func (l *TurnLoop[T, M]) Wait() *TurnLoopExitState[T, M] { return l.result } +func (l *TurnLoop[T, M]) takePendingResume(ctx context.Context) (*turnLoopPendingResume[T], bool) { + l.resumeMu.Lock() + pr := l.pendingResume + if pr == nil { + l.resumeMu.Unlock() + return nil, false + } + l.pendingResume = nil + l.resumeMu.Unlock() + return pr, true +} + +func (l *TurnLoop[T, M]) restorePendingResume(pr *turnLoopPendingResume[T]) { + if pr == nil { + return + } + l.resumeMu.Lock() + defer l.resumeMu.Unlock() + l.pendingResume = pr +} + +type turnLoopNextItems[T any] struct { + isResume bool + pr *turnLoopPendingResume[T] + items []T + pushBack []T +} + +func (l *TurnLoop[T, M]) collectNextTurnItems(ctx context.Context) (*turnLoopNextItems[T], bool) { + next := &turnLoopNextItems[T]{} + if l.pendingResume != nil { + next.isResume = true + var ok bool + next.pr, ok = l.takePendingResume(ctx) + if !ok { + return nil, false + } + + l.preemptCtrl.waitForPushes() + buffered := l.buffer.TakeAll() + if !next.pr.resumeSubmitted { + next.pr.resumeItems = append(next.pr.resumeItems, buffered...) + } else { + next.pr.unhandled = append(next.pr.unhandled, buffered...) + } + + next.pushBack = make([]T, 0, len(next.pr.interrupted)+len(next.pr.unhandled)+len(next.pr.resumeItems)) + next.pushBack = append(next.pushBack, next.pr.interrupted...) + next.pushBack = append(next.pushBack, next.pr.unhandled...) + next.pushBack = append(next.pushBack, next.pr.resumeItems...) + return next, true + } + + first, ok := l.receiveNextTurnItem(ctx) + if !ok { + return nil, false + } + + if err := ctx.Err(); err != nil { + l.buffer.PushFront([]T{first}) + l.runErr = err + return nil, false + } + + if l.stopCtrl.isCommitted() { + l.buffer.PushFront([]T{first}) + return nil, false + } + + l.preemptCtrl.waitForPushes() + rest := l.buffer.TakeAll() + next.items = append([]T{first}, rest...) + next.pushBack = next.items + return next, true +} + +func (l *TurnLoop[T, M]) receiveNextTurnItem(ctx context.Context) (T, bool) { + if idleFor := l.stopCtrl.idleDuration(); idleFor > 0 { + return l.receiveNextTurnItemUntilIdle(ctx, idleFor) + } + first, ok := l.buffer.Receive() + // Woken up by Stop(UntilIdleFor); re-enter loop to start the idle timer. + if !ok && l.stopCtrl.idleDuration() > 0 { + var zero T + return zero, false + } + if !ok { + if err := ctx.Err(); err != nil { + l.runErr = err + } + } + return first, ok +} + +func (l *TurnLoop[T, M]) receiveNextTurnItemUntilIdle(ctx context.Context, idleFor time.Duration) (T, bool) { + l.buffer.ClearWakeup() + idleTimer := time.NewTimer(idleFor) + cancelIdle := make(chan struct{}) + // When the idle timer fires, commitStop closes the buffer via buffer.Close(), + // which broadcasts to unblock the pending Receive() call below. + go func() { + select { + case <-idleTimer.C: + l.commitStop() + case <-cancelIdle: + } + }() + + first, ok := l.buffer.Receive() + + idleTimer.Stop() + close(cancelIdle) + + if !ok { + if err := ctx.Err(); err != nil { + l.runErr = err + } + if !l.buffer.IsClosed() { + var zero T + return zero, false + } + } + return first, ok +} + func (l *TurnLoop[T, M]) run(ctx context.Context) { defer l.cleanup(ctx) @@ -1590,85 +1807,17 @@ func (l *TurnLoop[T, M]) run(ctx context.Context) { return } - isResume := false - var pr *turnLoopPendingResume[T] - var items []T - var pushBack []T - - if l.pendingResume != nil { - isResume = true - pr = l.pendingResume - l.pendingResume = nil - - l.preemptCtrl.waitForPushes() - pr.newItems = append(pr.newItems, l.buffer.TakeAll()...) - - pushBack = make([]T, 0, len(pr.interrupted)+len(pr.unhandled)+len(pr.newItems)) - pushBack = append(pushBack, pr.interrupted...) - pushBack = append(pushBack, pr.unhandled...) - pushBack = append(pushBack, pr.newItems...) - } else { - var first T - var ok bool - - if idleFor := l.stopCtrl.idleDuration(); idleFor > 0 { - l.buffer.ClearWakeup() - idleTimer := time.NewTimer(idleFor) - cancelIdle := make(chan struct{}) - // When the idle timer fires, commitStop closes the buffer via - // buffer.Close(), which broadcasts to unblock the pending - // Receive() call below. - go func() { - select { - case <-idleTimer.C: - l.commitStop() - case <-cancelIdle: - } - }() - - first, ok = l.buffer.Receive() - - idleTimer.Stop() - close(cancelIdle) - - // A spurious wakeup can occur if Stop(UntilIdleFor) called - // buffer.Wakeup() after ClearWakeup() above but before - // Receive() entered its wait. In that case, Receive returns - // !ok from the woken flag, not from buffer closure. - // Re-enter the loop so the idle timer restarts cleanly. - if !ok && !l.buffer.IsClosed() { - continue - } - } else { - first, ok = l.buffer.Receive() - // Woken up by Stop(UntilIdleFor); re-enter loop to start the idle timer. - if !ok && l.stopCtrl.idleDuration() > 0 { - continue - } - } - - if !ok { - if err := ctx.Err(); err != nil { - l.runErr = err - } - return - } - - if err := ctx.Err(); err != nil { - l.buffer.PushFront([]T{first}) - l.runErr = err - return - } - - if l.stopCtrl.isCommitted() { - l.buffer.PushFront([]T{first}) - return + next, ok := l.collectNextTurnItems(ctx) + if !ok { + if l.stopCtrl.idleDuration() > 0 && !l.stopCtrl.isCommitted() && !l.buffer.IsClosed() && ctx.Err() == nil { + continue } + return + } - l.preemptCtrl.waitForPushes() - rest := l.buffer.TakeAll() - items = append([]T{first}, rest...) - pushBack = items + if next.isResume && l.stopCtrl.isCommitted() { + l.restorePendingResume(next.pr) + return } l.preemptCtrl.beginPlanningTurn() @@ -1676,11 +1825,11 @@ func (l *TurnLoop[T, M]) run(ctx context.Context) { l.preemptCtrl.abortPlanningTurn().ack() } - plan, err := l.planTurn(ctx, isResume, items, pr) + plan, err := l.planTurn(ctx, next.isResume, next.items, next.pr) if err != nil { abortPlanning() - if len(pushBack) > 0 { - l.buffer.PushFront(pushBack) + if len(next.pushBack) > 0 { + l.buffer.PushFront(next.pushBack) } l.runErr = err return @@ -1688,8 +1837,12 @@ func (l *TurnLoop[T, M]) run(ctx context.Context) { if l.stopCtrl.isCommitted() { abortPlanning() - if len(pushBack) > 0 { - l.buffer.PushFront(pushBack) + if next.isResume && plan.spec.isResume { + l.restorePendingResume(next.pr) + return + } + if len(next.pushBack) > 0 { + l.buffer.PushFront(next.pushBack) } return } @@ -1697,8 +1850,11 @@ func (l *TurnLoop[T, M]) run(ctx context.Context) { agent, err := l.config.PrepareAgent(plan.turnCtx, l, plan.spec.consumed) if err != nil { abortPlanning() - if len(pushBack) > 0 { - l.buffer.PushFront(pushBack) + if len(next.pushBack) > 0 { + l.buffer.PushFront(next.pushBack) + } + if next.isResume && !plan.spec.isResume { + l.loadCheckpointID = "" } l.runErr = err return @@ -1706,15 +1862,34 @@ func (l *TurnLoop[T, M]) run(ctx context.Context) { if l.stopCtrl.isCommitted() { abortPlanning() - if len(pushBack) > 0 { - l.buffer.PushFront(pushBack) + if next.isResume && plan.spec.isResume { + l.restorePendingResume(next.pr) + return + } + if len(next.pushBack) > 0 { + l.buffer.PushFront(next.pushBack) } return } + if next.isResume && !plan.spec.isResume && l.loadCheckpointID != "" { + checkpointID := l.loadCheckpointID + if err := l.deleteTurnLoopCheckpoint(ctx, checkpointID); err != nil { + abortPlanning() + if len(next.pushBack) > 0 { + l.buffer.PushFront(next.pushBack) + } + l.loadCheckpointID = "" + l.runErr = fmt.Errorf("failed to abandon checkpoint[%s] before fresh turn: %w", checkpointID, err) + return + } + l.loadCheckpointID = "" + } + l.buffer.PushFront(plan.remaining) runErr := l.runAgentAndHandleEvents(plan.turnCtx, agent, plan.spec) + runErr = l.deleteLoadedCheckpointAfterSuccessfulResume(ctx, runErr, plan.spec.isResume, l.interruptContexts != nil) if runErr != nil { // Set interruptedItems when a cancel or interrupt was captured from the @@ -1737,19 +1912,20 @@ func (l *TurnLoop[T, M]) run(ctx context.Context) { } func (l *TurnLoop[T, M]) setupBridgeStore(spec *turnRunSpec[T, M], runOpts []AgentRunOption) ([]AgentRunOption, *bridgeStore, error) { - store := l.config.Store - if store == nil && spec.isResume { - return nil, nil, fmt.Errorf("failed to resume agent: checkpoint store is nil") - } - if store == nil { + needsBridge := l.config.Store != nil || spec.isResume + if !needsBridge { return runOpts, nil, nil } - runOpts = append(runOpts, WithCheckPointID(bridgeCheckpointID)) + checkpointID := bridgeCheckpointID + if spec.resumeCheckpointID != "" { + checkpointID = spec.resumeCheckpointID + } + runOpts = append(runOpts, WithCheckPointID(checkpointID)) if spec.isResume { if len(spec.resumeBytes) == 0 { return nil, nil, fmt.Errorf("resume checkpoint is empty") } - return runOpts, newResumeBridgeStore(bridgeCheckpointID, spec.resumeBytes), nil + return runOpts, newResumeBridgeStore(checkpointID, spec.resumeBytes), nil } return runOpts, newBridgeStore(), nil } @@ -1814,6 +1990,7 @@ func (l *TurnLoop[T, M]) runAgentAndHandleEvents( l.interruptContexts = nil l.capturedCancelErr = nil l.checkPointRunnerBytes = nil + l.checkPointRunnerID = "" var iter *AsyncIterator[*TypedAgentEvent[M]] @@ -1822,7 +1999,6 @@ func (l *TurnLoop[T, M]) runAgentAndHandleEvents( l.preemptCtrl.abortPlanningTurn().ack() return err } - store := l.config.Store cancelOpt, agentCancelFunc := WithCancel() runOpts = append(runOpts, cancelOpt) @@ -1833,10 +2009,17 @@ func (l *TurnLoop[T, M]) runAgentAndHandleEvents( if spec.input != nil { enableStreaming = spec.input.EnableStreaming } + var runnerStore CheckPointStore + if ms != nil { + runnerStore = ms + } runner := NewTypedRunner(TypedRunnerConfig[M]{ EnableStreaming: enableStreaming, Agent: agent, - CheckPointStore: ms, + CheckPointStore: runnerStore, + SessionID: l.config.SessionID, + SessionStore: l.config.SessionStore, + SessionConfig: l.config.SessionConfig, }) preemptDone := make(chan struct{}) @@ -1859,9 +2042,9 @@ func (l *TurnLoop[T, M]) runAgentAndHandleEvents( if spec.isResume { var err error if spec.resumeParams != nil { - iter, err = runner.ResumeWithParams(ctx, bridgeCheckpointID, spec.resumeParams, runOpts...) + iter, err = runner.ResumeWithParams(ctx, spec.resumeCheckpointID, spec.resumeParams, runOpts...) } else { - iter, err = runner.Resume(ctx, bridgeCheckpointID, runOpts...) + iter, err = runner.Resume(ctx, spec.resumeCheckpointID, runOpts...) } if err != nil { return fmt.Errorf("failed to resume agent: %w", err) @@ -1919,18 +2102,20 @@ func (l *TurnLoop[T, M]) runAgentAndHandleEvents( go l.watchStop(done, agentCancelFunc, stoppedDone) finalizeCheckpoint := func() error { - if store != nil && ms != nil { - data, ok, err := ms.Get(ctx, bridgeCheckpointID) - if err != nil { - return fmt.Errorf("failed to read runner checkpoint: %w", err) - } + if ms != nil { + key, data, ok := ms.LastCheckpoint() if ok { + l.checkPointRunnerID = key l.checkPointRunnerBytes = append([]byte{}, data...) } } return nil } + finish := func(err error) error { + return err + } + // Wait for the turn to end. Three outcomes: // // done: Events fully handled (normal or error). If Stop() was @@ -1959,7 +2144,7 @@ func (l *TurnLoop[T, M]) runAgentAndHandleEvents( handleErr = err } } - return l.applyFrameworkCapturedError(handleErr) + return finish(l.applyFrameworkCapturedError(handleErr)) case <-preemptDone: <-done return nil @@ -1972,7 +2157,7 @@ func (l *TurnLoop[T, M]) runAgentAndHandleEvents( handleErr = err } } - return l.applyFrameworkCapturedError(handleErr) + return finish(l.applyFrameworkCapturedError(handleErr)) } } @@ -1995,12 +2180,26 @@ func (l *TurnLoop[T, M]) applyFrameworkCapturedError(handleErr error) error { return nil } +func interruptedItemsForExit[T any](items []T, pending *turnLoopPendingResume[T]) []T { + if pending != nil { + return pending.interrupted + } + return items +} + func (l *TurnLoop[T, M]) cleanup(ctx context.Context) { atomic.StoreInt32(&l.stopped, 1) unhandled := l.buffer.TakeAll() + l.resumeMu.Lock() + pending := l.pendingResume + l.resumeMu.Unlock() + if pending != nil { + unhandled = append(append([]T{}, pending.unhandled...), unhandled...) + } checkpointID := l.config.CheckpointID - isIdle := len(l.checkPointRunnerBytes) == 0 && len(unhandled) == 0 && len(l.interruptedItems) == 0 + hasPendingRunnerState := pending != nil && len(pending.resumeBytes) > 0 + isIdle := len(l.checkPointRunnerBytes) == 0 && !hasPendingRunnerState && len(unhandled) == 0 && len(l.interruptedItems) == 0 // Only save checkpoint when the loop exited due to an explicit Stop(), // a CancelError, or a business interrupt (InterruptError). @@ -2009,18 +2208,32 @@ func (l *TurnLoop[T, M]) cleanup(ctx context.Context) { exitCausedByStop := l.runErr == nil || errors.As(l.runErr, new(*CancelError)) || l.capturedCancelErr != nil businessInterrupt := errors.As(l.runErr, new(*InterruptError)) || l.interruptContexts != nil shouldSaveCheckpoint := l.config.Store != nil && checkpointID != "" && - ((l.stopCtrl.isCommitted() && exitCausedByStop) || businessInterrupt) && + ((l.stopCtrl.isCommitted() && exitCausedByStop) || businessInterrupt || (pending != nil && hasPendingRunnerState)) && !isIdle && !l.stopCtrl.skipCheckpointEnabled() var checkpointed bool var checkpointErr error if shouldSaveCheckpoint { + runnerCheckpointID := l.checkPointRunnerID + runnerCheckpoint := l.checkPointRunnerBytes + interruptedItems := l.interruptedItems + var resumeItems []T + if pending != nil { + runnerCheckpointID = pending.resumeCheckpointID + runnerCheckpoint = pending.resumeBytes + interruptedItems = pending.interrupted + if pending.resumeSubmitted { + resumeItems = append([]T{}, pending.resumeItems...) + } + } cp := &turnLoopCheckpoint[T]{ - RunnerCheckpoint: l.checkPointRunnerBytes, - HasRunnerState: len(l.checkPointRunnerBytes) > 0, - UnhandledItems: unhandled, - CanceledItems: l.interruptedItems, + RunnerCheckpointID: runnerCheckpointID, + RunnerCheckpoint: runnerCheckpoint, + HasRunnerState: len(runnerCheckpoint) > 0, + UnhandledItems: unhandled, + ResumeItems: resumeItems, + CanceledItems: interruptedItems, } checkpointed = true checkpointErr = l.saveTurnLoopCheckpoint(ctx, checkpointID, cp) @@ -2034,7 +2247,7 @@ func (l *TurnLoop[T, M]) cleanup(ctx context.Context) { l.result = &TurnLoopExitState[T, M]{ ExitReason: l.runErr, UnhandledItems: unhandled, - InterruptedItems: l.interruptedItems, + InterruptedItems: interruptedItemsForExit(l.interruptedItems, pending), StopCause: l.stopCtrl.cause(), CheckpointAttempted: checkpointed, CheckpointErr: checkpointErr, diff --git a/adk/turn_loop_test.go b/adk/turn_loop_test.go index 2fb903b20..79d3cc89a 100644 --- a/adk/turn_loop_test.go +++ b/adk/turn_loop_test.go @@ -28,6 +28,9 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" "github.com/cloudwego/eino/schema" ) @@ -87,8 +90,6 @@ type turnLoopCancellableMockAgent struct { name string runFunc func(ctx context.Context, input *AgentInput) (*AgentOutput, error) onCancel func(cc *cancelContext) - cancel context.CancelFunc - mu sync.Mutex } func (a *turnLoopCancellableMockAgent) Name(_ context.Context) string { return a.name } @@ -100,10 +101,8 @@ func (a *turnLoopCancellableMockAgent) Run(ctx context.Context, input *AgentInpu o := getCommonOptions(nil, opts...) cc := o.cancelCtx - a.mu.Lock() var cancelCtx context.Context - cancelCtx, a.cancel = context.WithCancel(ctx) - a.mu.Unlock() + cancelCtx, cancel := context.WithCancel(ctx) go func() { defer gen.Close() @@ -117,11 +116,7 @@ func (a *turnLoopCancellableMockAgent) Run(ctx context.Context, input *AgentInpu if a.onCancel != nil { a.onCancel(cc) } - a.mu.Lock() - if a.cancel != nil { - a.cancel() - } - a.mu.Unlock() + cancel() }() } @@ -981,7 +976,7 @@ func TestTurnLoop_GetAgentError_RecoverConsumed(t *testing.T) { func TestTurnLoop_GenInputError_RecoverItems(t *testing.T) { genErr := errors.New("gen input error") - loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { return nil, genErr }, @@ -990,6 +985,7 @@ func TestTurnLoop_GenInputError_RecoverItems(t *testing.T) { loop.Push("msg1") loop.Push("msg2") + loop.Run(context.Background()) result := loop.Wait() assert.ErrorIs(t, result.ExitReason, genErr) @@ -1846,6 +1842,88 @@ func TestTurnLoop_BusinessInterrupt_PersistAndResume(t *testing.T) { assert.Equal(t, []string{"msg1"}, resumeInterruptedItems, "interruptedItems should contain the original items") } +func TestTurnLoop_RestoredPendingResumeDistinguishesLegacyAndAcceptedResumeItems(t *testing.T) { + ctx := context.Background() + + run := func(t *testing.T, checkpoint *turnLoopCheckpoint[string], pushedBeforeRun string) (resumeItems []string, unhandledItems []string) { + t.Helper() + + store := newTestStore() + cpID := "restored-pending-" + pushedBeforeRun + data, err := marshalTurnLoopCheckpoint(checkpoint) + require.NoError(t, err) + require.NoError(t, store.Set(ctx, cpID, data)) + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], interruptedItems, unhandledItems, resumeItems []string) (*GenResumeResult[string, *schema.Message], error) { + return &GenResumeResult[string, *schema.Message]{ + Decision: TurnLoopResumeDecisionStartNewTurn, + Input: &AgentInput{}, + Consumed: interruptedItems, + }, nil + }, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + ok, ack := loop.Push(pushedBeforeRun) + require.True(t, ok) + require.Nil(t, ack) + + loop.config.GenResume = func(ctx context.Context, _ *TurnLoop[string, *schema.Message], interruptedItems, gotUnhandledItems, gotResumeItems []string) (*GenResumeResult[string, *schema.Message], error) { + unhandledItems = append([]string{}, gotUnhandledItems...) + resumeItems = append([]string{}, gotResumeItems...) + return &GenResumeResult[string, *schema.Message]{ + Decision: TurnLoopResumeDecisionStartNewTurn, + Input: &AgentInput{}, + Consumed: interruptedItems, + }, nil + } + + loop.Run(ctx) + exit := loop.Wait() + require.NoError(t, exit.ExitReason) + return resumeItems, unhandledItems + } + + t.Run("legacy restored checkpoint treats pre-run buffered item as resume intent", func(t *testing.T) { + resumeItems, unhandledItems := run(t, &turnLoopCheckpoint[string]{ + HasRunnerState: true, + RunnerCheckpointID: "runner-cp", + RunnerCheckpoint: []byte("runner-bytes"), + CanceledItems: []string{"interrupted"}, + UnhandledItems: []string{"normal-before-stop"}, + }, "legacy-resume") + + assert.Equal(t, []string{"legacy-resume"}, resumeItems) + assert.Equal(t, []string{"normal-before-stop"}, unhandledItems) + }) + + t.Run("persisted resume items keep pre-run buffered item as normal unhandled input", func(t *testing.T) { + resumeItems, unhandledItems := run(t, &turnLoopCheckpoint[string]{ + HasRunnerState: true, + RunnerCheckpointID: "runner-cp", + RunnerCheckpoint: []byte("runner-bytes"), + CanceledItems: []string{"interrupted"}, + UnhandledItems: []string{"normal-before-stop"}, + ResumeItems: []string{"accepted-resume"}, + }, "future-normal") + + assert.Equal(t, []string{"accepted-resume"}, resumeItems) + assert.Equal(t, []string{"normal-before-stop", "future-normal"}, unhandledItems) + }) +} + // turnLoopInterruptAgent is a test agent that produces a business interrupt event. type turnLoopInterruptAgent struct { interruptInfo any @@ -2086,6 +2164,7 @@ type deletableCheckpointStore struct { turnLoopCheckpointStore deleteCalled bool deletedKey string + deleteErr error } func (s *deletableCheckpointStore) Delete(_ context.Context, key string) error { @@ -2093,6 +2172,9 @@ func (s *deletableCheckpointStore) Delete(_ context.Context, key string) error { defer s.mu.Unlock() s.deleteCalled = true s.deletedKey = key + if s.deleteErr != nil { + return s.deleteErr + } delete(s.m, key) return nil } @@ -2367,7 +2449,7 @@ func TestTurnLoop_ResumeWaitsForInFlightPushBeforePlanning(t *testing.T) { }) loop.pendingResume = &turnLoopPendingResume[string]{ interrupted: []string{"interrupted"}, - newItems: []string{"pre-existing"}, + resumeItems: []string{"pre-existing"}, } go func() { @@ -2805,6 +2887,235 @@ func TestNewTurnLoop_WaitBeforeRun(t *testing.T) { } } +type mockSessionStore struct { + mu sync.Mutex + events map[string][]storedSessionEvent +} + +func (m *mockSessionStore) AppendEvents(ctx context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + return m.AppendEventsForSession(ctx, sessionID, events) +} + +func (m *mockSessionStore) AppendEventsForSession(_ context.Context, sessionID string, events []*SessionEvent[*schema.Message]) error { + m.mu.Lock() + defer m.mu.Unlock() + if m.events == nil { + m.events = make(map[string][]storedSessionEvent) + } + for _, event := range events { + if event == nil || event.EventID == "" { + return ErrInvalidEventID + } + if err := NormalizeSessionEventKind(event); err != nil { + return err + } + data, err := encodeSessionEvent(event) + if err != nil { + return err + } + m.events[sessionID] = append(m.events[sessionID], storedSessionEvent{ + EventID: event.EventID, + Kind: event.Kind, + Data: data, + }) + } + return nil +} + +func (m *mockSessionStore) LoadEvents(ctx context.Context, sessionID string, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + return m.LoadEventsForSession(ctx, sessionID, req) +} + +func (m *mockSessionStore) LoadEventsForSession(_ context.Context, sessionID string, opts *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + m.mu.Lock() + defer m.mu.Unlock() + if opts == nil { + opts = &LoadSessionEventsRequest{} + } + events := m.events[sessionID] + findAfter := func() (int, error) { + if opts.After == "" { + return -1, nil + } + for i, event := range events { + if event.EventID == opts.After { + return i, nil + } + } + return -1, ErrEventIDOutOfRange + } + after, err := findAfter() + if err != nil { + return nil, err + } + kindSet := buildTestKindSet(opts.Kinds) + var out []*SessionEvent[*schema.Message] + hasMore := false + if opts.Reverse { + end := len(events) + if opts.After != "" { + end = after + } + for i := end - 1; i >= 0; i-- { + if kindSet != nil { + if _, ok := kindSet[events[i].Kind]; !ok { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + event, err := decodeSessionEvent[*schema.Message](events[i].Data) + if err != nil { + return nil, err + } + out = append(out, event) + } + } else { + for i := after + 1; i < len(events); i++ { + if kindSet != nil { + if _, ok := kindSet[events[i].Kind]; !ok { + continue + } + } + if opts.Limit > 0 && len(out) >= opts.Limit { + hasMore = true + break + } + event, err := decodeSessionEvent[*schema.Message](events[i].Data) + if err != nil { + return nil, err + } + out = append(out, event) + } + } + var next string + if hasMore && len(out) > 0 { + next = out[len(out)-1].EventID + } + return &LoadSessionEventsResult[*schema.Message]{Events: out, Next: next}, nil +} + +func (m *mockSessionStore) openSession(_ context.Context, req *openSessionRequest) (*openSessionResult[*schema.Message], error) { + sessionID := "" + if req != nil { + sessionID = req.sessionID + } + return &openSessionResult[*schema.Message]{ + handle: &mockSessionHandle{store: m, sessionID: sessionID}, + }, nil +} + +type mockSessionHandle struct { + store *mockSessionStore + sessionID string +} + +func (h *mockSessionHandle) loadEvents(ctx context.Context, req *LoadSessionEventsRequest) (*LoadSessionEventsResult[*schema.Message], error) { + if req == nil { + req = &LoadSessionEventsRequest{} + } + return h.store.LoadEventsForSession(ctx, h.sessionID, req) +} + +func (h *mockSessionHandle) appendEvents(ctx context.Context, events []*SessionEvent[*schema.Message]) error { + return h.store.AppendEventsForSession(ctx, h.sessionID, events) +} + +func (h *mockSessionHandle) close(context.Context) error { return nil } + +func TestTurnLoop_SessionStoreWithCheckpointIDWithoutStore(t *testing.T) { + ctx := context.Background() + sessionID := "test-session-id" + sessionStore := &mockSessionStore{} + var processed bool + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeFirst, + PrepareAgent: func(context.Context, *TurnLoop[string, *schema.Message], []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(context.Context, *AgentInput) (*AgentOutput, error) { + processed = true + return &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage("response", nil), + Role: schema.Assistant, + }, + }, nil + }, + }, nil + }, + OnAgentEvents: func(_ context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + SessionID: sessionID, + SessionStore: sessionStore, + CheckpointID: "test-checkpoint-id", + }) + + loop.Push("test-message") + loop.Run(ctx) + exit := loop.Wait() + + assert.NoError(t, exit.ExitReason) + assert.True(t, processed) + assert.NotEmpty(t, sessionStore.events[sessionID]) +} + +func TestTurnLoop_SessionStoreWithoutCheckpointStoreSkipsRunnerCheckpoint(t *testing.T) { + ctx := context.Background() + sessionID := "test-session-without-checkpoint-store" + sessionStore := &mockSessionStore{} + var processed bool + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeFirst, + PrepareAgent: func(context.Context, *TurnLoop[string, *schema.Message], []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(context.Context, *AgentInput) (*AgentOutput, error) { + processed = true + return &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage("response", nil), + Role: schema.Assistant, + }, + }, nil + }, + }, nil + }, + OnAgentEvents: func(_ context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + SessionID: sessionID, + SessionStore: sessionStore, + }) + + loop.Push("test-message") + loop.Run(ctx) + exit := loop.Wait() + + assert.NoError(t, exit.ExitReason) + assert.True(t, processed) + assert.NotEmpty(t, sessionStore.events[sessionID]) +} + func TestNewTurnLoop_RunIsIdempotent(t *testing.T) { var genInputCalls int32 @@ -5440,7 +5751,7 @@ func TestStopController_CloseForLoopExitClearsPendingCancel(t *testing.T) { assert.False(t, ok) } -func TestAttack_UntilIdleFor_ConcurrentPushDuringIdleTimer(t *testing.T) { +func TestTurnLoop_UntilIdleFor_ConcurrentPushDuringIdleTimer(t *testing.T) { turnCount := int32(0) turnDone := make(chan struct{}, 10) @@ -5481,7 +5792,7 @@ func TestAttack_UntilIdleFor_ConcurrentPushDuringIdleTimer(t *testing.T) { assert.Equal(t, int32(6), finalCount, "all 6 pushes should have been processed") } -func TestAttack_UntilIdleFor_MultipleStopCallsFirstWins(t *testing.T) { +func TestTurnLoop_UntilIdleFor_MultipleStopCallsFirstWins(t *testing.T) { turnDone := make(chan struct{}) loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ GenInput: genInputConsumeAllWithMsg, @@ -5511,7 +5822,7 @@ func TestAttack_UntilIdleFor_MultipleStopCallsFirstWins(t *testing.T) { waitOrFail(t, done, "second UntilIdleFor should have been ignored; loop should have exited with 100ms timer") } -func TestAttack_BareStopOverridesUntilIdleFor(t *testing.T) { +func TestTurnLoop_Stop_BareStopOverridesUntilIdleFor(t *testing.T) { agentStarted := make(chan struct{}) agentDone := make(chan struct{}) @@ -5549,7 +5860,7 @@ func TestAttack_BareStopOverridesUntilIdleFor(t *testing.T) { assert.NoError(t, exit.ExitReason, "bare Stop should exit cleanly") } -func TestAttack_BareStopDoesNotDeescalateExistingCancelIntent(t *testing.T) { +func TestTurnLoop_Stop_BareStopDoesNotDeescalateExistingCancelIntent(t *testing.T) { agentStarted := make(chan *cancelContext, 1) probe := &turnLoopStopModeProbeAgent{ccCh: agentStarted} loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ @@ -5578,7 +5889,7 @@ func TestAttack_BareStopDoesNotDeescalateExistingCancelIntent(t *testing.T) { assert.Equal(t, CancelImmediate, ce.Info.Mode) } -func TestAttack_InterruptedItems_EmptyWhenAgentFinishesNormally(t *testing.T) { +func TestTurnLoop_InterruptedItems_EmptyWhenAgentFinishesNormally(t *testing.T) { agentStarted := make(chan struct{}) loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ GenInput: genInputConsumeAllWithMsg, @@ -5603,7 +5914,7 @@ func TestAttack_InterruptedItems_EmptyWhenAgentFinishesNormally(t *testing.T) { assert.Empty(t, exit.InterruptedItems, "InterruptedItems must be empty when agent finished normally") } -func TestAttack_TurnBuffer_WakeupDoesNotLoseItems(t *testing.T) { +func TestTurnBuffer_WakeupDoesNotLoseItems(t *testing.T) { tb := newTurnBuffer[string]() tb.Send("a") @@ -5621,7 +5932,7 @@ func TestAttack_TurnBuffer_WakeupDoesNotLoseItems(t *testing.T) { assert.Equal(t, []string{"a", "b", "c"}, got, "Wakeup must not cause items to be lost") } -func TestAttack_TurnBuffer_ClearWakeupPreventsSpuriousReturn(t *testing.T) { +func TestTurnBuffer_ClearWakeupPreventsSpuriousReturn(t *testing.T) { tb := newTurnBuffer[string]() tb.Wakeup() @@ -5646,7 +5957,7 @@ func TestAttack_TurnBuffer_ClearWakeupPreventsSpuriousReturn(t *testing.T) { } } -func TestAttack_StopBeforeRun_UntilIdleFor_ExitsImmediately(t *testing.T) { +func TestTurnLoop_StopBeforeRun_UntilIdleForExitsImmediately(t *testing.T) { loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ GenInput: genInputConsumeAllWithMsg, PrepareAgent: prepareTestAgent, @@ -5666,7 +5977,7 @@ func TestAttack_StopBeforeRun_UntilIdleFor_ExitsImmediately(t *testing.T) { waitOrFail(t, done, "loop should exit immediately when Stop() called before Run()") } -func TestAttack_PushAfterStop_UntilIdleFor_RoutedToLateItems(t *testing.T) { +func TestTurnLoop_PushAfterStop_UntilIdleForRoutedToLateItems(t *testing.T) { turnDone := make(chan struct{}) loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ GenInput: genInputConsumeAllWithMsg, @@ -5695,7 +6006,7 @@ func TestAttack_PushAfterStop_UntilIdleFor_RoutedToLateItems(t *testing.T) { assert.Equal(t, []string{"after-stop"}, late) } -func TestAttack_ConcurrentStopEscalation_RaceDetector(t *testing.T) { +func TestTurnLoop_Stop_ConcurrentEscalation(t *testing.T) { agentStarted := make(chan struct{}) loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ GenInput: genInputConsumeAllWithMsg, @@ -5737,7 +6048,7 @@ func TestAttack_ConcurrentStopEscalation_RaceDetector(t *testing.T) { t.Log("ExitReason:", exit.ExitReason) } -func TestAttack_SkipCheckpoint_Sticky(t *testing.T) { +func TestTurnLoop_Stop_SkipCheckpointSticky(t *testing.T) { agentStarted := make(chan struct{}) loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ GenInput: genInputConsumeAllWithMsg, @@ -5897,10 +6208,15 @@ func TestSaveTurnLoopCheckpoint_NilStore(t *testing.T) { func TestSetupBridgeStore_NilStore_Resume(t *testing.T) { l := &TurnLoop[string, *schema.Message]{config: TurnLoopConfig[string, *schema.Message]{Store: nil}} - spec := &turnRunSpec[string, *schema.Message]{isResume: true} - _, _, err := l.setupBridgeStore(spec, nil) - assert.Error(t, err) - assert.Contains(t, err.Error(), "checkpoint store is nil") + spec := &turnRunSpec[string, *schema.Message]{isResume: true, resumeCheckpointID: "runner-cp", resumeBytes: []byte("runner-bytes")} + opts, ms, err := l.setupBridgeStore(spec, nil) + require.NoError(t, err) + require.NotNil(t, ms) + assert.Len(t, opts, 1) + data, ok, err := ms.Get(context.Background(), "runner-cp") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, []byte("runner-bytes"), data) } // TestTurnLoop_Preempt_LoopStalledAfterSecondPreemptPush covers a liveness @@ -5976,7 +6292,7 @@ func TestTurnLoop_Preempt_LoopStalledAfterSecondPreemptPush(t *testing.T) { assert.Equal(t, int32(3), atomic.LoadInt32(&turnCount), "expected 3 turns to be processed") } -func TestAttack_BusinessInterrupt_NoStore_ExitsWithoutPanic(t *testing.T) { +func TestTurnLoop_BusinessInterrupt_NoStoreExitsWithoutPanic(t *testing.T) { ctx := context.Background() interruptAgent := &turnLoopInterruptAgent{interruptInfo: "no_store_test"} @@ -5996,7 +6312,7 @@ func TestAttack_BusinessInterrupt_NoStore_ExitsWithoutPanic(t *testing.T) { assert.False(t, exit.CheckpointAttempted, "no store → no checkpoint attempt") } -func TestAttack_BusinessInterrupt_EmptyConsumed_NoCheckpoint(t *testing.T) { +func TestTurnLoop_BusinessInterrupt_EmptyConsumedNoCheckpoint(t *testing.T) { ctx := context.Background() store := newTestStore() interruptAgent := &turnLoopInterruptAgent{interruptInfo: "idle_test"} @@ -6022,3 +6338,225 @@ func TestAttack_BusinessInterrupt_EmptyConsumed_NoCheckpoint(t *testing.T) { require.True(t, errors.As(exit.ExitReason, &intErr), "expected *InterruptError, got: %v", exit.ExitReason) assert.Empty(t, exit.InterruptedItems, "consumed was empty → InterruptedItems should be empty") } + +func TestTurnLoop_BusinessInterrupt_WithStorePersistsCheckpoint(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "business-interrupt-with-store" + interruptAgent := &turnLoopInterruptAgent{interruptInfo: "approval_needed"} + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return interruptAgent, nil + }, + }) + + loop.Push("msg1") + exit := loop.Wait() + + var intErr *InterruptError + require.True(t, errors.As(exit.ExitReason, &intErr), "expected *InterruptError, got: %v", exit.ExitReason) + assert.Equal(t, []string{"msg1"}, exit.InterruptedItems) + require.True(t, exit.CheckpointAttempted, "store is configured → checkpoint should be attempted") + require.NoError(t, exit.CheckpointErr) + + store.mu.Lock() + data, ok := store.m[cpID] + store.mu.Unlock() + require.True(t, ok, "checkpoint should exist in store") + + cp, err := unmarshalTurnLoopCheckpoint[string](data) + require.NoError(t, err) + assert.True(t, cp.HasRunnerState, "checkpoint should have runner state") + assert.NotEmpty(t, cp.RunnerCheckpoint, "checkpoint should have runner checkpoint bytes") + assert.Equal(t, []string{"msg1"}, cp.CanceledItems, "checkpoint should have canceled items") +} + +func TestTurnLoop_BusinessInterrupt_WithoutStoreExitsCleanly(t *testing.T) { + ctx := context.Background() + interruptAgent := &turnLoopInterruptAgent{interruptInfo: "no_store_clean"} + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return interruptAgent, nil + }, + }) + + loop.Push("msg1") + exit := loop.Wait() + + var intErr *InterruptError + require.True(t, errors.As(exit.ExitReason, &intErr), "expected *InterruptError, got: %v", exit.ExitReason) + assert.Equal(t, []string{"msg1"}, exit.InterruptedItems) + assert.False(t, exit.CheckpointAttempted, "no store → no checkpoint attempt") + assert.Nil(t, loop.pendingResume, "no managed pending resume should exist") +} + +type turnLoopAgenticToolCallModel struct { + callCount int32 +} + +func (m *turnLoopAgenticToolCallModel) Generate(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + if atomic.AddInt32(&m.callCount, 1) == 1 { + return agenticToolCallMsg("turn_loop_slow_tool", "call-1", `{"input":"x"}`), nil + } + return agenticMsg("done"), nil +} + +func (m *turnLoopAgenticToolCallModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.AgenticMessage{msg}), nil +} + +func TestTurnLoop_StopGracefulThenImmediate_AgenticStreamableToolCheckpoint(t *testing.T) { + ctx := context.Background() + + gate := make(chan struct{}) + slowTool := &slowStreamingTool{ + name: "turn_loop_slow_tool", + chunkInterval: time.Hour, + chunks: []string{"chunk"}, + started: make(chan struct{}, 1), + gate: gate, + } + t.Cleanup(func() { + close(gate) + }) + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TurnLoopAgenticRepro", + Description: "repro agent", + Model: &turnLoopAgenticToolCallModel{}, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{slowTool}}, + }, + }) + require.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.AgenticMessage]{ + Store: newTestStore(), + GenInput: func(_ context.Context, _ *TurnLoop[string, *schema.AgenticMessage], items []string) (*GenInputResult[string, *schema.AgenticMessage], error) { + return &GenInputResult[string, *schema.AgenticMessage]{ + Input: &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage(items[0])}, + }, + Consumed: items, + }, nil + }, + PrepareAgent: func(_ context.Context, _ *TurnLoop[string, *schema.AgenticMessage], _ []string) (TypedAgent[*schema.AgenticMessage], error) { + return agent, nil + }, + }) + + loop.Push("trigger") + select { + case <-slowTool.started: + case <-time.After(5 * time.Second): + t.Fatal("streamable tool did not start") + } + + loop.Stop(WithGraceful()) + time.Sleep(50 * time.Millisecond) + loop.Stop(WithImmediate()) + + exit := loop.Wait() + + var cancelErr *CancelError + require.True(t, errors.As(exit.ExitReason, &cancelErr), "ExitReason should be a *CancelError, got %v", exit.ExitReason) + assert.NoError(t, exit.CheckpointErr) +} + +func TestTurnLoop_PreemptAfterToolCallsTimeout_AgenticStreamableToolCheckpoint(t *testing.T) { + ctx := context.Background() + + gate := make(chan struct{}) + slowTool := &slowStreamingTool{ + name: "turn_loop_slow_tool", + chunkInterval: time.Millisecond, + chunks: []string{"chunk-1", "chunk-2", "chunk-3"}, + started: make(chan struct{}, 1), + gate: gate, + } + t.Cleanup(func() { + close(gate) + }) + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TurnLoopAgenticPreemptRepro", + Description: "repro agent", + Model: &turnLoopAgenticToolCallModel{}, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{slowTool}}, + }, + }) + require.NoError(t, err) + + errCh := make(chan error, 16) + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.AgenticMessage]{ + Store: newTestStore(), + GenInput: func(_ context.Context, _ *TurnLoop[string, *schema.AgenticMessage], items []string) (*GenInputResult[string, *schema.AgenticMessage], error) { + return &GenInputResult[string, *schema.AgenticMessage]{ + Input: &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage(items[0])}, + }, + Consumed: []string{items[0]}, + Remaining: func() []string { + if len(items) <= 1 { + return nil + } + return append([]string(nil), items[1:]...) + }(), + }, nil + }, + PrepareAgent: func(_ context.Context, _ *TurnLoop[string, *schema.AgenticMessage], _ []string) (TypedAgent[*schema.AgenticMessage], error) { + return agent, nil + }, + OnAgentEvents: func(_ context.Context, _ *TurnContext[string, *schema.AgenticMessage], events *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) error { + for { + ev, ok := events.Next() + if !ok { + return nil + } + if ev.Err != nil { + errCh <- ev.Err + } + } + }, + }) + + loop.Push("trigger") + select { + case <-slowTool.started: + case <-time.After(5 * time.Second): + t.Fatal("streamable tool did not start") + } + time.Sleep(20 * time.Millisecond) + + ok, ack := loop.Push("preempt", WithPreemptTimeout[string, *schema.AgenticMessage](AfterToolCalls, 20*time.Millisecond)) + require.True(t, ok) + select { + case <-ack: + case <-time.After(5 * time.Second): + t.Fatal("preempt was not acknowledged") + } + + loop.Stop() + exit := loop.Wait() + + for { + select { + case err := <-errCh: + assert.NotContains(t, err.Error(), "gob marshal error") + default: + assert.NoError(t, exit.CheckpointErr) + return + } + } +} diff --git a/adk/usage.go b/adk/usage.go new file mode 100644 index 000000000..269a3b49f --- /dev/null +++ b/adk/usage.go @@ -0,0 +1,72 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import "github.com/cloudwego/eino/schema" + +func assistantTokenUsage[M MessageType](msg M) *schema.TokenUsage { + switch v := any(msg).(type) { + case *schema.Message: + if v == nil || v.Role != schema.Assistant || v.ResponseMeta == nil { + return nil + } + return v.ResponseMeta.Usage + case *schema.AgenticMessage: + if v == nil || v.Role != schema.AgenticRoleTypeAssistant || v.ResponseMeta == nil { + return nil + } + return v.ResponseMeta.TokenUsage + default: + return nil + } +} + +func assistantFinishReason[M MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + if v == nil || v.Role != schema.Assistant || v.ResponseMeta == nil { + return "" + } + return v.ResponseMeta.FinishReason + case *schema.AgenticMessage: + if v == nil || v.Role != schema.AgenticRoleTypeAssistant || v.ResponseMeta == nil { + return "" + } + if v.ResponseMeta.ClaudeExtension != nil { + return v.ResponseMeta.ClaudeExtension.StopReason + } + if v.ResponseMeta.GeminiExtension != nil { + return v.ResponseMeta.GeminiExtension.FinishReason + } + return "" + default: + return "" + } +} + +func modelUsageFromAssistant[M MessageType](msg M) *ModelUsage { + usage := assistantTokenUsage(msg) + if usage == nil { + return nil + } + return &ModelUsage{ + InputTokens: usage.PromptTokens, + OutputTokens: usage.CompletionTokens, + CacheReadInputTokens: usage.PromptTokenDetails.CachedTokens, + Raw: usage, + } +} diff --git a/adk/wrappers.go b/adk/wrappers.go index ef5c03537..3c20d40f8 100644 --- a/adk/wrappers.go +++ b/adk/wrappers.go @@ -22,6 +22,7 @@ import ( "io" "reflect" "sync" + "time" "github.com/google/uuid" @@ -294,12 +295,277 @@ type typedEventSenderModel[M MessageType] struct { modelFailoverConfig *ModelFailoverConfig[M] } +func sendSessionTimelineEvent[M MessageType](ctx context.Context, se *SessionEvent[M]) { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil || execCtx.generator == nil || se == nil { + return + } + if !execCtx.timelineEvents && !execCtx.internalTimelineEvents { + return + } + if se.Timestamp.IsZero() { + se.Timestamp = newEventTimestamp() + } + if se.EventID == "" { + if err := assignSessionEventIDFromContext(ctx, se); err != nil { + execCtx.send(ctx, &TypedAgentEvent[M]{Err: err}) + return + } + } + if err := ValidateEmittedSessionEventKind(se); err != nil { + execCtx.send(ctx, &TypedAgentEvent[M]{Err: err}) + return + } + execCtx.send(ctx, &TypedAgentEvent[M]{SessionEventVariant: &SessionEventVariant[M]{Event: se}}) +} + +func newModelSpanStartEvent[M MessageType](ctx context.Context, spanID string, started time.Time, opts ...model.Option) *SessionEvent[M] { + meta := modelSpanMetaFromContext[M](ctx, opts...) + meta.Model.Accepted = false + return &SessionEvent[M]{ + Timestamp: started, + Kind: SessionEventSpanModelRequestStart, + Span: &SpanEvent{ + SpanID: spanID, + Kind: SpanKindModel, + Name: "model_request", + StartedAt: started, + ParentSpanID: meta.ParentSpanID, + Model: meta.Model, + }, + } +} + +type modelSpanEndEventInput[M MessageType] struct { + spanID string + startEventID string + started time.Time + ended time.Time + msg M + err error + accepted bool + firstChunk time.Duration +} + +func newModelSpanEndEvent[M MessageType](ctx context.Context, in modelSpanEndEventInput[M], opts ...model.Option) *SessionEvent[M] { + status := "ok" + errStr := "" + if in.err != nil { + status = "error" + errStr = in.err.Error() + if errors.Is(in.err, context.Canceled) || errors.Is(in.err, ErrStreamCanceled) { + status = "cancelled" + } + } + return &SessionEvent[M]{ + Timestamp: in.ended, + Kind: SessionEventSpanModelRequestEnd, + Span: &SpanEvent{ + SpanID: in.spanID, + Kind: SpanKindModel, + Name: "model_request", + StartedAt: in.started, + EndedAt: in.ended, + TTFTMS: in.firstChunk.Milliseconds(), + Status: status, + Err: errStr, + ParentSpanID: modelSpanMetaFromContext[M](ctx, opts...).ParentSpanID, + Model: modelSpanCompletionMeta(ctx, in.startEventID, in.msg, in.accepted && in.err == nil, in.err, opts...), + }, + } +} + +type modelSpanContextMeta struct { + ParentSpanID string + Model *ModelSpanMeta +} + +func modelSpanMetaFromContext[M MessageType](ctx context.Context, opts ...model.Option) modelSpanContextMeta { + meta := &ModelSpanMeta{Attempt: 1} + if common := model.GetCommonOptions(nil, opts...); common != nil && common.Model != nil { + meta.Model = *common.Model + } + if currentModel, ok := typedGetFailoverCurrentModel[M](ctx); ok { + if provider, ok := components.GetType(currentModel); ok { + meta.Provider = provider + } + } + parentSpanID := "" + if failoverMeta, ok := getFailoverTimeline(ctx); ok { + parentSpanID = failoverMeta.ParentSpanID + if failoverMeta.Attempt > 0 { + meta.Attempt = failoverMeta.Attempt + } + } else { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + meta.Attempt = st.getRetryAttempt() + 1 + return nil + }) + } + return modelSpanContextMeta{ParentSpanID: parentSpanID, Model: meta} +} + +func modelSpanCompletionMeta[M MessageType](ctx context.Context, startEventID string, msg M, accepted bool, err error, opts ...model.Option) *ModelSpanMeta { + meta := modelSpanMetaFromContext[M](ctx, opts...).Model + meta.ModelRequestStartEventID = startEventID + meta.Usage = modelUsageFromAssistant(msg) + meta.FinishReason = assistantFinishReason(msg) + meta.Accepted = accepted + var timeoutErr interface { + ModelTimeoutSpanMeta() (phase string, timeout time.Duration, elapsed time.Duration, chunksReceived int) + } + if errors.As(err, &timeoutErr) { + phase, timeout, elapsed, chunksReceived := timeoutErr.ModelTimeoutSpanMeta() + meta.Timeout = &ModelTimeoutMeta{ + Phase: phase, + TimeoutMS: timeout.Milliseconds(), + ElapsedMS: elapsed.Milliseconds(), + ChunksReceived: chunksReceived, + } + } + return meta +} + +// lookupOrInitToolSpanInFlight returns the existing in-flight entry for the +// given tCtx.CallID (signalling a resumed call), or initializes a fresh entry +// (snapshotting CurrentModelSpanID / CurrentAssistantMessageEventID from +// typedState). It does NOT yet write the entry into typedState — the caller +// is responsible for invoking persistToolSpanInFlight after emitting the +// tool_call_start span and capturing its EventID. +// +// Reads (and the conditional snapshot) happen inside compose.ProcessState so +// that concurrent parallel tool calls are serialized safely — the framework +// guarantees the closure runs with exclusive access to typedState. +func lookupOrInitToolSpanInFlight[M MessageType](ctx context.Context, tCtx *ToolContext) (*toolSpanInFlight, bool) { + var ( + entry *toolSpanInFlight + isResume bool + ) + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + if existing, ok := st.ToolSpansInFlight[tCtx.CallID]; ok && existing != nil { + entry = existing + isResume = true + return nil + } + entry = &toolSpanInFlight{ + SpanID: uuid.NewString(), + StartedAt: newEventTimestamp(), + ParentSpanID: st.CurrentModelSpanID, + AssistantMessageEventID: st.CurrentAssistantMessageEventID, + } + return nil + }) + return entry, isResume +} + +// persistToolSpanInFlight writes the in-flight entry into typedState.ToolSpansInFlight +// keyed by callID. Called from the start-emission path after the start +// SessionEvent's EventID has been captured into the entry. +func persistToolSpanInFlight[M MessageType](ctx context.Context, callID string, entry *toolSpanInFlight) { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + if st.ToolSpansInFlight == nil { + st.ToolSpansInFlight = make(map[string]*toolSpanInFlight) + } + st.ToolSpansInFlight[callID] = entry + return nil + }) +} + +// clearToolSpanInFlight removes the in-flight entry for callID. Called after +// the matching tool_call_end span has been emitted. +func clearToolSpanInFlight[M MessageType](ctx context.Context, callID string) { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + if st.ToolSpansInFlight == nil { + return nil + } + delete(st.ToolSpansInFlight, callID) + return nil + }) +} + +func newToolSpanStartEvent[M MessageType](ctx context.Context, inFlight *toolSpanInFlight, tCtx *ToolContext) *SessionEvent[M] { + return &SessionEvent[M]{ + Timestamp: inFlight.StartedAt, + Kind: SessionEventSpanToolCallStart, + Span: &SpanEvent{ + SpanID: inFlight.SpanID, + ParentSpanID: inFlight.ParentSpanID, + Kind: SpanKindTool, + Name: "tool_call", + StartedAt: inFlight.StartedAt, + Tool: &ToolSpanMeta{ + ToolUseID: tCtx.CallID, + Name: tCtx.Name, + AssistantMessageEventID: inFlight.AssistantMessageEventID, + }, + }, + } +} + +type toolSpanEndEventInput struct { + ended time.Time + err error + resultEventID string +} + +func newToolSpanEndEvent[M MessageType](ctx context.Context, inFlight *toolSpanInFlight, tCtx *ToolContext, in toolSpanEndEventInput) *SessionEvent[M] { + status := "ok" + errStr := "" + if in.err != nil { + status = "error" + errStr = in.err.Error() + if errors.Is(in.err, context.Canceled) || errors.Is(in.err, ErrStreamCanceled) { + status = "cancelled" + } + } + ended := in.ended + if ended.IsZero() { + ended = newEventTimestamp() + } + return &SessionEvent[M]{ + Timestamp: ended, + Kind: SessionEventSpanToolCallEnd, + Span: &SpanEvent{ + SpanID: inFlight.SpanID, + ParentSpanID: inFlight.ParentSpanID, + Kind: SpanKindTool, + Name: "tool_call", + StartedAt: inFlight.StartedAt, + EndedAt: ended, + Status: status, + Err: errStr, + Tool: &ToolSpanMeta{ + ToolUseID: tCtx.CallID, + Name: tCtx.Name, + ToolCallStartEventID: inFlight.StartEventID, + AssistantMessageEventID: inFlight.AssistantMessageEventID, + ToolResultMessageEventID: in.resultEventID, + }, + }, + } +} + func (m *typedEventSenderModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + started := newEventTimestamp() + spanID := uuid.NewString() + startEvent := newModelSpanStartEvent[M](ctx, spanID, started, opts...) + sendSessionTimelineEvent(ctx, startEvent) result, err := m.inner.Generate(ctx, input, opts...) + ended := newEventTimestamp() + sendSessionTimelineEvent(ctx, newModelSpanEndEvent(ctx, modelSpanEndEventInput[M]{ + spanID: spanID, + startEventID: startEvent.EventID, + started: started, + ended: ended, + msg: result, + err: err, + accepted: err == nil, + }, opts...)) if err != nil { var zero M return zero, err } + timestamp := newEventTimestamp() execCtx := getTypedChatModelAgentExecCtx[M](ctx) if execCtx != nil && execCtx.suppressEventSend { @@ -310,17 +576,54 @@ func (m *typedEventSenderModel[M]) Generate(ctx context.Context, input []M, opts return zero, errors.New("generator is nil when sending event in Generate: ensure agent state is properly initialized") } + // Build a SessionEventMessage draft for the assistant message and route + // its ID allocation through the runner's SessionEventIDGenerator[M] so + // producer-owned identity applies. The same ID is used for the live + // TypedAgentEvent below; the materialized SessionEvent later inherits it. + assistantDraft := &SessionEvent[M]{Timestamp: timestamp, Kind: SessionEventMessage, Message: copyMessage(result)} + if err := assignSessionEventIDFromContext(ctx, assistantDraft); err != nil { + var zero M + return zero, err + } + assistantMsgEventID := assistantDraft.EventID + + // Persist the model span ID and assistant message event ID into typedState + // so the tool wrapper can snapshot them into per-call ToolSpansInFlight + // entries when emitting tool_call_start spans. The snapshot survives + // interrupt/resume; the matching tool_call_end span (which may fire on + // a later run) reads ParentSpanID and AssistantMessageEventID from the + // snapshot, preserving the link to the original turn's model output. + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.CurrentModelSpanID = spanID + st.CurrentAssistantMessageEventID = assistantMsgEventID + return nil + }) + event := typedModelOutputEvent(copyMessage(result), nil) - execCtx.send(event) + event.SessionEventVariant = &SessionEventVariant[M]{Event: assistantDraft} + execCtx.send(ctx, event) return result, nil } func (m *typedEventSenderModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + started := newEventTimestamp() + spanID := uuid.NewString() + startEvent := newModelSpanStartEvent[M](ctx, spanID, started, opts...) + sendSessionTimelineEvent(ctx, startEvent) result, err := m.inner.Stream(ctx, input, opts...) if err != nil { + sendSessionTimelineEvent(ctx, newModelSpanEndEvent(ctx, modelSpanEndEventInput[M]{ + spanID: spanID, + startEventID: startEvent.EventID, + started: started, + ended: newEventTimestamp(), + msg: *new(M), + err: err, + }, opts...)) return nil, err } + timestamp := newEventTimestamp() execCtx := getTypedChatModelAgentExecCtx[M](ctx) if execCtx == nil || execCtx.generator == nil { @@ -328,7 +631,7 @@ func (m *typedEventSenderModel[M]) Stream(ctx context.Context, input []M, opts . return nil, errors.New("generator is nil when sending event in Stream: ensure agent state is properly initialized") } - streams := result.Copy(2) + streams := result.Copy(3) eventStream := streams[0] if convertOpts := m.buildStreamConvertOptions(ctx); len(convertOpts) > 0 { @@ -337,13 +640,112 @@ func (m *typedEventSenderModel[M]) Stream(ctx context.Context, input []M, opts . convertOpts...) } + // Build a streaming-mode draft for the assistant message; the message + // itself is materialized later by the consumer, but we route ID + // allocation through the runner's SessionEventIDGenerator[M] now so the + // live TypedAgentEvent and the eventual SessionEvent share a producer- + // owned ID. Generators that need to recognize the assistant draft can + // match on Kind==SessionEventMessage with a zero Message. + var draftZero M + assistantDraft := &SessionEvent[M]{Timestamp: timestamp, Kind: SessionEventMessage, Message: draftZero} + if err := assignSessionEventIDFromContext(ctx, assistantDraft); err != nil { + result.Close() + return nil, err + } + assistantMsgEventID := assistantDraft.EventID + + // Persist the model span ID and assistant message event ID into typedState + // so the tool wrapper can snapshot them into per-call ToolSpansInFlight + // entries when emitting tool_call_start spans. The snapshot survives + // interrupt/resume; the matching tool_call_end span (which may fire on + // a later run) reads ParentSpanID and AssistantMessageEventID from the + // snapshot, preserving the link to the original turn's model output. + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.CurrentModelSpanID = spanID + st.CurrentAssistantMessageEventID = assistantMsgEventID + return nil + }) + var zero M event := typedModelOutputEvent[M](zero, eventStream) - execCtx.send(event) + event.SessionEventVariant = &SessionEventVariant[M]{ + MessageStreamRef: &MessageStreamRef{ + EventID: assistantMsgEventID, + Timestamp: timestamp, + Kind: SessionEventMessage, + }, + } + execCtx.send(ctx, event) + + spanStream := streams[2] + go func() { + firstChunk := time.Duration(0) + firstAt := time.Time{} + var chunks []M + var streamErr error + for { + msg, recvErr := spanStream.Recv() + if recvErr == io.EOF { + break + } + if recvErr != nil { + streamErr = recvErr + break + } + if firstAt.IsZero() { + firstAt = newEventTimestamp() + firstChunk = firstAt.Sub(started) + } + chunks = append(chunks, msg) + } + spanStream.Close() + var final M + if len(chunks) > 0 && streamErr == nil { + final, streamErr = concatMessagesForSpan(chunks) + } + sendSessionTimelineEvent(ctx, newModelSpanEndEvent(ctx, modelSpanEndEventInput[M]{ + spanID: spanID, + startEventID: startEvent.EventID, + started: started, + ended: newEventTimestamp(), + msg: final, + err: streamErr, + accepted: streamErr == nil, + firstChunk: firstChunk, + }, opts...)) + }() return streams[1], nil } +func concatMessagesForSpan[M MessageType](chunks []M) (M, error) { + var zero M + switch any(zero).(type) { + case *schema.Message: + msgs := make([]*schema.Message, 0, len(chunks)) + for _, chunk := range chunks { + msgs = append(msgs, any(chunk).(*schema.Message)) + } + msg, err := schema.ConcatMessages(msgs) + if err != nil { + return zero, err + } + return any(msg).(M), nil + case *schema.AgenticMessage: + msgs := make([]*schema.AgenticMessage, 0, len(chunks)) + for _, chunk := range chunks { + msgs = append(msgs, any(chunk).(*schema.AgenticMessage)) + } + msg, err := schema.ConcatAgenticMessages(msgs) + if err != nil { + return zero, err + } + return any(msg).(M), nil + default: + return zero, nil + } +} + // buildStreamConvertOptions constructs ConvertOption hooks that gate stream termination behind // the retry verdict signal protocol. // @@ -506,13 +908,56 @@ func GetMessageID[M MessageType](msg M) string { // EnsureMessageID assigns a UUID v4 message ID if the message doesn't have one. // Idempotent: if ID already set, no-op. -// Middleware authors should call this before SendEvent if they create messages. +// TypedSendEvent/SendEvent call this automatically for message-bearing events. +// Middleware authors only need to call it directly when they need the ID before +// emitting the event. func EnsureMessageID[M MessageType](msg M) { switch v := any(msg).(type) { case *schema.Message: - v.Extra = internal.EnsureMessageID(v.Extra) + if internal.GetMessageID(v.Extra) == "" { + v.Extra = internal.EnsureMessageID(v.Extra) + } case *schema.AgenticMessage: - v.Extra = internal.EnsureMessageID(v.Extra) + if internal.GetMessageID(v.Extra) == "" { + v.Extra = internal.EnsureMessageID(v.Extra) + } + } +} + +func ensureTypedAgentEventMessageIDs[M MessageType](event *TypedAgentEvent[M]) { + if event == nil { + return + } + if event.Output != nil && event.Output.MessageOutput != nil && !isNilMessage(event.Output.MessageOutput.Message) { + EnsureMessageID(event.Output.MessageOutput.Message) + } + if event.SessionEventVariant != nil { + ensureSessionEventMessageIDs(event.SessionEventVariant.Event) + } +} + +func ensureSessionEventMessageIDs[M MessageType](event *SessionEvent[M]) { + if event == nil { + return + } + if !isNilMessage(event.Message) { + EnsureMessageID(event.Message) + } + if event.MessagesReplaced != nil { + for _, msg := range *event.MessagesReplaced { + if !isNilMessage(msg) { + EnsureMessageID(msg) + } + } + } + if event.MessageUpdated != nil && !isNilMessage(event.MessageUpdated.Message) { + msgID := GetMessageID(event.MessageUpdated.Message) + if msgID == "" && event.MessageUpdated.MessageID != "" { + typedSetMessageID(event.MessageUpdated.Message, event.MessageUpdated.MessageID) + } + } + if event.MessageInserted != nil && !isNilMessage(event.MessageInserted.Message) { + EnsureMessageID(event.MessageInserted.Message) } } @@ -843,10 +1288,32 @@ func typedToolEnhancedStreamEvent[M MessageType](callID, toolName, toolMsgID str func (w *typedEventSenderToolWrapper[M]) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error) { return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + inFlight, isResume := lookupOrInitToolSpanInFlight[M](ctx, tCtx) + if !isResume { + startEvent := newToolSpanStartEvent[M](ctx, inFlight, tCtx) + sendSessionTimelineEvent(ctx, startEvent) + inFlight.StartEventID = startEvent.EventID + persistToolSpanInFlight[M](ctx, tCtx.CallID, inFlight) + } + result, err := endpoint(ctx, argumentsInJSON, opts...) if err != nil { + if _, isInterrupt := compose.IsInterruptRerunError(err); isInterrupt { + // An interrupt-shape error means the tool did not complete; the call is + // paused awaiting resume. Defer the end span: leave the in-flight entry + // in typedState so the next invocation of this wrapper for the same + // CallID reuses the SpanID and emits the matching end span. See §3.1 + // and §3.6 of the design plan for the full lifecycle. + return "", err + } + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, toolSpanEndEventInput{ + ended: newEventTimestamp(), + err: err, + })) + clearToolSpanInFlight[M](ctx, tCtx.CallID) return "", err } + timestamp := newEventTimestamp() toolName := tCtx.Name callID := tCtx.CallID @@ -854,6 +1321,22 @@ func (w *typedEventSenderToolWrapper[M]) WrapInvokableToolCall(_ context.Context prePopAction := typedPopToolGenAction[M](ctx, toolName) toolMsgID := uuid.NewString() event := typedToolInvokeEvent[M](callID, toolName, result, toolMsgID) + // Route the tool result message ID through the runner's + // SessionEventIDGenerator[M] via a SessionEventMessage draft so + // custom-tool-result generators see the populated message. Fail-closed: + // on allocation failure, skip both the tool result event and the + // matching tool span end so no orphaned ToolResultMessageEventID + // reference is left in the timeline. + toolResultDraft := &SessionEvent[M]{Timestamp: timestamp, Kind: SessionEventMessage, Message: event.Output.MessageOutput.Message} + if idErr := assignSessionEventIDFromContext(ctx, toolResultDraft); idErr != nil { + if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil && execCtx.generator != nil { + execCtx.send(ctx, &TypedAgentEvent[M]{Err: idErr}) + } + clearToolSpanInFlight[M](ctx, tCtx.CallID) + return "", idErr + } + resultEventID := toolResultDraft.EventID + event.SessionEventVariant = &SessionEventVariant[M]{Event: toolResultDraft} if prePopAction != nil { event.Action = prePopAction } @@ -864,21 +1347,45 @@ func (w *typedEventSenderToolWrapper[M]) WrapInvokableToolCall(_ context.Context if st.getReturnDirectlyToolCallID() == callID { st.setReturnDirectlyEvent(event) } else { - execCtx.send(event) + execCtx.send(ctx, event) } return nil }) + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, toolSpanEndEventInput{ + ended: newEventTimestamp(), + resultEventID: resultEventID, + })) + clearToolSpanInFlight[M](ctx, tCtx.CallID) + return result, nil }, nil } func (w *typedEventSenderToolWrapper[M]) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error) { return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + inFlight, isResume := lookupOrInitToolSpanInFlight[M](ctx, tCtx) + if !isResume { + startEvent := newToolSpanStartEvent[M](ctx, inFlight, tCtx) + sendSessionTimelineEvent(ctx, startEvent) + inFlight.StartEventID = startEvent.EventID + persistToolSpanInFlight[M](ctx, tCtx.CallID, inFlight) + } + result, err := endpoint(ctx, argumentsInJSON, opts...) if err != nil { + if _, isInterrupt := compose.IsInterruptRerunError(err); isInterrupt { + // Defer end span; in-flight entry remains for resume. See §3.1. + return nil, err + } + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, toolSpanEndEventInput{ + ended: newEventTimestamp(), + err: err, + })) + clearToolSpanInFlight[M](ctx, tCtx.CallID) return nil, err } + timestamp := newEventTimestamp() toolName := tCtx.Name callID := tCtx.CallID @@ -887,7 +1394,78 @@ func (w *typedEventSenderToolWrapper[M]) WrapStreamableToolCall(_ context.Contex streams := result.Copy(2) toolMsgID := uuid.NewString() + // Streaming tool result: the materialized message is not yet + // available, so the draft carries a zero M and SessionEventMessage + // kind. ID allocation flows through the runner's + // SessionEventIDGenerator[M]. Fail-closed: on allocation failure, + // skip the tool result event AND the tool span end so no orphaned + // ToolResultMessageEventID reference is left behind. + var toolResultDraftMsg M + toolResultDraft := &SessionEvent[M]{Timestamp: timestamp, Kind: SessionEventMessage, Message: toolResultDraftMsg} + if idErr := assignSessionEventIDFromContext(ctx, toolResultDraft); idErr != nil { + if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil && execCtx.generator != nil { + execCtx.send(ctx, &TypedAgentEvent[M]{Err: idErr}) + } + streams[0].Close() + streams[1].Close() + clearToolSpanInFlight[M](ctx, tCtx.CallID) + return nil, idErr + } + resultEventID := toolResultDraft.EventID + + // End-span emission for streamable tools attaches to the caller's + // stream copy via schema.WithOnEOF (success path) and + // schema.WithErrWrapper (error / cancellation path). Both hooks fire + // synchronously inside the consumer's recv() call, so the end span + // is enqueued before the agent's event generator closes — this is + // what avoids the race the previous goroutine drainer hit. + // + // Residual risk: the hooks fire only when the consumer drives the + // stream to a terminal state (io.EOF or non-EOF error). If a + // consumer abandons the stream mid-flight (e.g. calls Close() early, + // or a tool implementation ignores ctx and produces unbounded + // chunks while the consumer stops calling Recv), neither hook fires + // and only the start span is persisted. Correctness is unaffected; + // observability shows an unmatched start span. If this ever becomes + // load-bearing, the fix is to reintroduce a goroutine drainer as a + // fallback gated by spanEndOnce, with a per-run WaitGroup on the + // exec ctx so the agent waits for it before closing the generator. + // + // emitEnd is attached to the caller's stream copy via WithOnEOF and + // WithErrWrapper. v3 makes it interrupt-aware: an interrupt-shape + // streamErr means the tool did not complete on this run, so neither + // emit the end span nor delete the in-flight entry. The next resume + // re-invokes this wrapper, sees the in-flight entry, and continues + // until a terminal state (success EOF, hard error, or cancellation) + // fires. See §3.5 of the design plan for the full lifecycle. + var spanEndOnce sync.Once + emitEnd := func(streamErr error) { + spanEndOnce.Do(func() { + if streamErr != nil { + if _, isInterrupt := compose.IsInterruptRerunError(streamErr); isInterrupt { + return + } + } + in := toolSpanEndEventInput{ + ended: newEventTimestamp(), + err: streamErr, + } + if streamErr == nil { + in.resultEventID = resultEventID + } + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, in)) + clearToolSpanInFlight[M](ctx, tCtx.CallID) + }) + } + event := typedToolStreamEvent[M](callID, toolName, toolMsgID, streams[0]) + event.SessionEventVariant = &SessionEventVariant[M]{ + MessageStreamRef: &MessageStreamRef{ + EventID: resultEventID, + Timestamp: timestamp, + Kind: SessionEventMessage, + }, + } event.Action = prePopAction execCtx := getTypedChatModelAgentExecCtx[M](ctx) @@ -896,21 +1474,50 @@ func (w *typedEventSenderToolWrapper[M]) WrapStreamableToolCall(_ context.Contex if st.getReturnDirectlyToolCallID() == callID { st.setReturnDirectlyEvent(event) } else { - execCtx.send(event) + execCtx.send(ctx, event) } return nil }) - return streams[1], nil + callerStream := schema.StreamReaderWithConvert(streams[1], + func(s string) (string, error) { return s, nil }, + schema.WithOnEOF(func() (any, error) { + emitEnd(nil) + return nil, io.EOF + }), + schema.WithErrWrapper(func(streamErr error) error { + emitEnd(streamErr) + return streamErr + }), + ) + return callerStream, nil }, nil } func (w *typedEventSenderToolWrapper[M]) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error) { return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + inFlight, isResume := lookupOrInitToolSpanInFlight[M](ctx, tCtx) + if !isResume { + startEvent := newToolSpanStartEvent[M](ctx, inFlight, tCtx) + sendSessionTimelineEvent(ctx, startEvent) + inFlight.StartEventID = startEvent.EventID + persistToolSpanInFlight[M](ctx, tCtx.CallID, inFlight) + } + result, err := endpoint(ctx, toolArgument, opts...) if err != nil { + if _, isInterrupt := compose.IsInterruptRerunError(err); isInterrupt { + // Defer end span; in-flight entry remains for resume. See §3.1. + return nil, err + } + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, toolSpanEndEventInput{ + ended: newEventTimestamp(), + err: err, + })) + clearToolSpanInFlight[M](ctx, tCtx.CallID) return nil, err } + timestamp := newEventTimestamp() toolName := tCtx.Name callID := tCtx.CallID @@ -919,8 +1526,28 @@ func (w *typedEventSenderToolWrapper[M]) WrapEnhancedInvokableToolCall(_ context toolMsgID := uuid.NewString() event, eventErr := typedToolEnhancedInvokeEvent[M](callID, toolName, toolMsgID, result) if eventErr != nil { + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, toolSpanEndEventInput{ + ended: newEventTimestamp(), + err: eventErr, + })) + clearToolSpanInFlight[M](ctx, tCtx.CallID) return nil, eventErr } + // Route the enhanced-invoke tool result message ID through the + // runner's SessionEventIDGenerator[M] via a SessionEventMessage + // draft. Fail-closed: on allocation failure, skip both the tool + // result event and the matching tool span end so no orphaned + // ToolResultMessageEventID reference is left behind. + toolResultDraft := &SessionEvent[M]{Timestamp: timestamp, Kind: SessionEventMessage, Message: event.Output.MessageOutput.Message} + if idErr := assignSessionEventIDFromContext(ctx, toolResultDraft); idErr != nil { + if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil && execCtx.generator != nil { + execCtx.send(ctx, &TypedAgentEvent[M]{Err: idErr}) + } + clearToolSpanInFlight[M](ctx, tCtx.CallID) + return nil, idErr + } + resultEventID := toolResultDraft.EventID + event.SessionEventVariant = &SessionEventVariant[M]{Event: toolResultDraft} if prePopAction != nil { event.Action = prePopAction } @@ -931,21 +1558,45 @@ func (w *typedEventSenderToolWrapper[M]) WrapEnhancedInvokableToolCall(_ context if st.getReturnDirectlyToolCallID() == callID { st.setReturnDirectlyEvent(event) } else { - execCtx.send(event) + execCtx.send(ctx, event) } return nil }) + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, toolSpanEndEventInput{ + ended: newEventTimestamp(), + resultEventID: resultEventID, + })) + clearToolSpanInFlight[M](ctx, tCtx.CallID) + return result, nil }, nil } func (w *typedEventSenderToolWrapper[M]) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error) { return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + inFlight, isResume := lookupOrInitToolSpanInFlight[M](ctx, tCtx) + if !isResume { + startEvent := newToolSpanStartEvent[M](ctx, inFlight, tCtx) + sendSessionTimelineEvent(ctx, startEvent) + inFlight.StartEventID = startEvent.EventID + persistToolSpanInFlight[M](ctx, tCtx.CallID, inFlight) + } + result, err := endpoint(ctx, toolArgument, opts...) if err != nil { + if _, isInterrupt := compose.IsInterruptRerunError(err); isInterrupt { + // Defer end span; in-flight entry remains for resume. See §3.1. + return nil, err + } + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, toolSpanEndEventInput{ + ended: newEventTimestamp(), + err: err, + })) + clearToolSpanInFlight[M](ctx, tCtx.CallID) return nil, err } + timestamp := newEventTimestamp() toolName := tCtx.Name callID := tCtx.CallID @@ -954,7 +1605,78 @@ func (w *typedEventSenderToolWrapper[M]) WrapEnhancedStreamableToolCall(_ contex streams := result.Copy(2) toolMsgID := uuid.NewString() + // Streaming tool result: the materialized message is not yet + // available, so the draft carries a zero M and SessionEventMessage + // kind. ID allocation flows through the runner's + // SessionEventIDGenerator[M]. Fail-closed: on allocation failure, + // skip the tool result event AND the tool span end so no orphaned + // ToolResultMessageEventID reference is left behind. + var toolResultDraftMsg M + toolResultDraft := &SessionEvent[M]{Timestamp: timestamp, Kind: SessionEventMessage, Message: toolResultDraftMsg} + if idErr := assignSessionEventIDFromContext(ctx, toolResultDraft); idErr != nil { + if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil && execCtx.generator != nil { + execCtx.send(ctx, &TypedAgentEvent[M]{Err: idErr}) + } + streams[0].Close() + streams[1].Close() + clearToolSpanInFlight[M](ctx, tCtx.CallID) + return nil, idErr + } + resultEventID := toolResultDraft.EventID + + // End-span emission for streamable tools attaches to the caller's + // stream copy via schema.WithOnEOF (success path) and + // schema.WithErrWrapper (error / cancellation path). Both hooks fire + // synchronously inside the consumer's recv() call, so the end span + // is enqueued before the agent's event generator closes — this is + // what avoids the race the previous goroutine drainer hit. + // + // Residual risk: the hooks fire only when the consumer drives the + // stream to a terminal state (io.EOF or non-EOF error). If a + // consumer abandons the stream mid-flight (e.g. calls Close() early, + // or a tool implementation ignores ctx and produces unbounded + // chunks while the consumer stops calling Recv), neither hook fires + // and only the start span is persisted. Correctness is unaffected; + // observability shows an unmatched start span. If this ever becomes + // load-bearing, the fix is to reintroduce a goroutine drainer as a + // fallback gated by spanEndOnce, with a per-run WaitGroup on the + // exec ctx so the agent waits for it before closing the generator. + // + // emitEnd is attached to the caller's stream copy via WithOnEOF and + // WithErrWrapper. v3 makes it interrupt-aware: an interrupt-shape + // streamErr means the tool did not complete on this run, so neither + // emit the end span nor delete the in-flight entry. The next resume + // re-invokes this wrapper, sees the in-flight entry, and continues + // until a terminal state (success EOF, hard error, or cancellation) + // fires. See §3.5 of the design plan for the full lifecycle. + var spanEndOnce sync.Once + emitEnd := func(streamErr error) { + spanEndOnce.Do(func() { + if streamErr != nil { + if _, isInterrupt := compose.IsInterruptRerunError(streamErr); isInterrupt { + return + } + } + in := toolSpanEndEventInput{ + ended: newEventTimestamp(), + err: streamErr, + } + if streamErr == nil { + in.resultEventID = resultEventID + } + sendSessionTimelineEvent(ctx, newToolSpanEndEvent[M](ctx, inFlight, tCtx, in)) + clearToolSpanInFlight[M](ctx, tCtx.CallID) + }) + } + event := typedToolEnhancedStreamEvent[M](callID, toolName, toolMsgID, streams[0]) + event.SessionEventVariant = &SessionEventVariant[M]{ + MessageStreamRef: &MessageStreamRef{ + EventID: resultEventID, + Timestamp: timestamp, + Kind: SessionEventMessage, + }, + } event.Action = prePopAction execCtx := getTypedChatModelAgentExecCtx[M](ctx) @@ -963,15 +1685,33 @@ func (w *typedEventSenderToolWrapper[M]) WrapEnhancedStreamableToolCall(_ contex if st.getReturnDirectlyToolCallID() == callID { st.setReturnDirectlyEvent(event) } else { - execCtx.send(event) + execCtx.send(ctx, event) } return nil }) - return streams[1], nil + callerStream := schema.StreamReaderWithConvert(streams[1], + func(tr *schema.ToolResult) (*schema.ToolResult, error) { return tr, nil }, + schema.WithOnEOF(func() (any, error) { + emitEnd(nil) + return nil, io.EOF + }), + schema.WithErrWrapper(func(streamErr error) error { + emitEnd(streamErr) + return streamErr + }), + ) + return callerStream, nil }, nil } +// drainStringToolResultForSpan and drainEnhancedToolResultForSpan are no +// longer needed; the streamable wrappers attach end-span emission via +// schema.WithOnEOF / schema.WithErrWrapper hooks on the caller's stream copy +// instead. The hook approach fires synchronously during the consumer's read +// loop, eliminating the goroutine race where the agent's event generator +// could close before the drainer's emission landed. + func hasUserEventSenderToolWrapper[M MessageType](handlers []TypedChatModelAgentMiddleware[M]) bool { for _, handler := range handlers { if _, ok := any(handler).(eventSenderToolWrapperMarker); ok { @@ -1204,6 +1944,9 @@ func (w *typedStateModelWrapper[M]) Generate(ctx context.Context, _ []M, opts .. } else { stateToolInfos = w.toolInfos } + if stateDeferredToolInfos == nil && composeLevelOpts.DeferredTools != nil { + stateDeferredToolInfos = composeLevelOpts.DeferredTools + } } state := &TypedChatModelAgentState[M]{ @@ -1242,6 +1985,7 @@ func (w *typedStateModelWrapper[M]) Generate(ctx context.Context, _ []M, opts .. st.DeferredToolInfos = state.DeferredToolInfos return nil }) + syncModelContextSessionEvent(ctx, state) // Derive model options from state. Append after caller opts so state takes precedence // (model.GetCommonOptions applies left-to-right, last wins). @@ -1270,6 +2014,7 @@ func (w *typedStateModelWrapper[M]) Generate(ctx context.Context, _ []M, opts .. }) } + EnsureMessageID(result) state.Messages = append(state.Messages, result) for _, handler := range w.handlers { @@ -1329,6 +2074,9 @@ func (w *typedStateModelWrapper[M]) Stream(ctx context.Context, _ []M, opts ...m } else { stateToolInfos = w.toolInfos } + if stateDeferredToolInfos == nil && composeLevelOpts.DeferredTools != nil { + stateDeferredToolInfos = composeLevelOpts.DeferredTools + } } state := &TypedChatModelAgentState[M]{ @@ -1365,6 +2113,7 @@ func (w *typedStateModelWrapper[M]) Stream(ctx context.Context, _ []M, opts ...m st.DeferredToolInfos = state.DeferredToolInfos return nil }) + syncModelContextSessionEvent(ctx, state) // Derive model options from state. Append after caller opts so state takes precedence // (model.GetCommonOptions applies left-to-right, last wins). @@ -1394,6 +2143,7 @@ func (w *typedStateModelWrapper[M]) Stream(ctx context.Context, _ []M, opts ...m }) } + EnsureMessageID(result) state.Messages = append(state.Messages, result) for _, handler := range w.handlers { diff --git a/adk/wrappers_failover_test.go b/adk/wrappers_failover_test.go deleted file mode 100644 index 45fb0c222..000000000 --- a/adk/wrappers_failover_test.go +++ /dev/null @@ -1,215 +0,0 @@ -/* - * Copyright 2026 CloudWeGo Authors - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package adk - -import ( - "context" - "errors" - "sync/atomic" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/cloudwego/eino/components/model" - "github.com/cloudwego/eino/schema" -) - -func TestBuildModelWrappers_FailoverProxyInner(t *testing.T) { - base := &fakeChatModel{ - callbacksEnabled: true, - generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - return schema.AssistantMessage("ok", nil), nil - }, - stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil - }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 0, - ShouldFailover: func(context.Context, *schema.Message, error) bool { return false }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - return base, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](base, &modelWrapperConfig{ - failoverConfig: failoverCfg, - }) - - smw, ok := wrapped.(*stateModelWrapper) - require.True(t, ok) - _, ok = smw.inner.(*failoverProxyModel) - require.True(t, ok) - require.Same(t, base, smw.original) - require.Same(t, failoverCfg, smw.modelFailoverConfig) -} - -func TestStateModelWrapper_Generate_WithFailover(t *testing.T) { - wantErr := errors.New("first failed") - var shouldCalls int32 - var m1Calls int32 - var m2Calls int32 - - m1 := &fakeChatModel{ - callbacksEnabled: true, - generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m1Calls, 1) - return schema.AssistantMessage("partial", nil), wantErr - }, - stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - return nil, errors.New("unused") - }, - } - m2 := &fakeChatModel{ - callbacksEnabled: true, - generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m2Calls, 1) - return schema.AssistantMessage("ok", nil), nil - }, - stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - return nil, errors.New("unused") - }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool { - atomic.AddInt32(&shouldCalls, 1) - require.ErrorIs(t, err, wantErr) - require.NotNil(t, out) - require.Equal(t, "partial", out.Content) - return true - }, - GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - require.Equal(t, uint(1), failoverCtx.FailoverAttempt) - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - got, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - require.NotNil(t, got) - require.Equal(t, "ok", got.Content) - require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) -} - -func TestStateModelWrapper_Stream_WithFailover(t *testing.T) { - streamErr := errors.New("mid error") - var shouldCalls int32 - var m1Calls int32 - var m2Calls int32 - - m1 := &fakeChatModel{ - callbacksEnabled: true, - generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - return nil, errors.New("unused") - }, - stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m1Calls, 1) - return streamWithMidError([]*schema.Message{ - schema.AssistantMessage("p1", nil), - schema.AssistantMessage("p2", nil), - }, streamErr), nil - }, - } - m2 := &fakeChatModel{ - callbacksEnabled: true, - generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - return nil, errors.New("unused") - }, - stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m2Calls, 1) - return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final", nil)}), nil - }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool { - atomic.AddInt32(&shouldCalls, 1) - require.ErrorIs(t, err, streamErr) - require.NotNil(t, out) - require.Equal(t, "p1p2", out.Content) - return true - }, - GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - require.Equal(t, uint(1), failoverCtx.FailoverAttempt) - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - msgs, err := drainMessageStream(sr) - require.NoError(t, err) - require.Len(t, msgs, 1) - require.Equal(t, "final", msgs[0].Content) - require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) -} - -func TestFailoverAcceptsAgenticAgent(t *testing.T) { - ctx := context.Background() - - m := &mockAgenticModel{ - generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { - return agenticMsg("ok"), nil - }, - } - - fallbackModel := &mockAgenticModel{ - generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { - return agenticMsg("fallback"), nil - }, - } - - agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ - Name: "FailoverAgent", - Description: "Agent with failover config", - Model: m, - ModelFailoverConfig: &ModelFailoverConfig[*schema.AgenticMessage]{ - MaxRetries: 1, - ShouldFailover: func(ctx context.Context, outputMessage *schema.AgenticMessage, outputErr error) bool { - return true - }, - GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext[*schema.AgenticMessage]) (model.BaseModel[*schema.AgenticMessage], []*schema.AgenticMessage, error) { - return fallbackModel, nil, nil - }, - }, - }) - require.NoError(t, err) - assert.NotNil(t, agent) -} diff --git a/adk/wrappers_retry_failover_test.go b/adk/wrappers_retry_failover_test.go deleted file mode 100644 index c1a291df6..000000000 --- a/adk/wrappers_retry_failover_test.go +++ /dev/null @@ -1,613 +0,0 @@ -/* - * Copyright 2026 CloudWeGo Authors - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package adk - -import ( - "context" - "errors" - "sync/atomic" - "testing" - "time" - - "github.com/stretchr/testify/require" - - "github.com/cloudwego/eino/components/model" - "github.com/cloudwego/eino/schema" -) - -func newFakeChatModel( - gen func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error), - stream func(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error), -) *fakeChatModel { - if gen == nil { - gen = func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { - return nil, errors.New("unused") - } - } - if stream == nil { - stream = func(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { - return nil, errors.New("unused") - } - } - return &fakeChatModel{callbacksEnabled: true, generate: gen, stream: stream} -} - -func TestRetryThenFailover(t *testing.T) { - t.Run("Generate_RetryExhaustedTriggersFailover", func(t *testing.T) { - modelErr := errors.New("model error") - var m1Calls int32 - var m2Calls int32 - - m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m1Calls, 1) - return nil, modelErr - }, nil) - m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m2Calls, 1) - return schema.AssistantMessage("ok from m2", nil), nil - }, nil) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 2, - IsRetryAble: func(_ context.Context, err error) bool { return true }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, fc *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - require.NotNil(t, fc.LastErr) - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - require.Equal(t, "ok from m2", msg.Content) - - // m1: 1 (lastSuccess) + 2 retries = 3 calls on lastSuccess attempt, - // then failover to m2 which also goes through retry wrapper: 1 call succeeds. - require.Equal(t, int32(3), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) - }) - - t.Run("Generate_AllExhausted", func(t *testing.T) { - modelErr := errors.New("always fails") - var m1Calls int32 - var m2Calls int32 - - m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m1Calls, 1) - return nil, modelErr - }, nil) - m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m2Calls, 1) - return nil, modelErr - }, nil) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 1, - IsRetryAble: func(_ context.Context, err error) bool { return true }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.Error(t, err) - - // Should be RetryExhaustedError from m2's retry wrapper - var retryErr *RetryExhaustedError - require.True(t, errors.As(err, &retryErr)) - - // m1: 1 initial + 1 retry = 2 calls - require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) - // m2: 1 initial + 1 retry = 2 calls - require.Equal(t, int32(2), atomic.LoadInt32(&m2Calls)) - }) - - t.Run("Generate_RetrySucceedsNoFailover", func(t *testing.T) { - var m1Calls int32 - var failoverCalled int32 - - m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - n := atomic.AddInt32(&m1Calls, 1) - if n == 1 { - return nil, errors.New("transient error") - } - return schema.AssistantMessage("ok on retry", nil), nil - }, nil) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 2, - IsRetryAble: func(_ context.Context, err error) bool { return true }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - atomic.AddInt32(&failoverCalled, 1) - return true - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - t.Fatal("GetFailoverModel should not be called when retry succeeds") - return nil, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - require.Equal(t, "ok on retry", msg.Content) - - // 2 calls: first fails, second succeeds via retry - require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) - // ShouldFailover should never be called - require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) - }) - - t.Run("Generate_NonRetryableErrorTriggersFailover", func(t *testing.T) { - nonRetryableErr := errors.New("non-retryable") - var m1Calls int32 - var m2Calls int32 - - m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m1Calls, 1) - return nil, nonRetryableErr - }, nil) - m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m2Calls, 1) - return schema.AssistantMessage("ok from m2", nil), nil - }, nil) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 3, - IsRetryAble: func(_ context.Context, err error) bool { - // Only non-retryable errors - return !errors.Is(err, nonRetryableErr) - }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - require.Equal(t, "ok from m2", msg.Content) - - // m1 called only once — non-retryable error skips retry - require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) - }) - - t.Run("Stream_RetryExhaustedTriggersFailover", func(t *testing.T) { - streamErr := errors.New("stream mid error") - var m1Calls int32 - var m2Calls int32 - - m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m1Calls, 1) - return streamWithMidError([]*schema.Message{ - schema.AssistantMessage("partial", nil), - }, streamErr), nil - }) - m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m2Calls, 1) - return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok from m2", nil)}), nil - }) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 1, - IsRetryAble: func(_ context.Context, err error) bool { return true }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, fc *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - require.NotNil(t, fc.LastErr) - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - msgs, err := drainMessageStream(sr) - require.NoError(t, err) - require.Len(t, msgs, 1) - require.Equal(t, "ok from m2", msgs[0].Content) - - // m1: 1 initial + 1 retry = 2 calls on lastSuccess attempt - require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) - }) - - t.Run("Stream_AllExhausted", func(t *testing.T) { - streamErr := errors.New("always fails mid-stream") - var m1Calls int32 - var m2Calls int32 - - m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m1Calls, 1) - return streamWithMidError([]*schema.Message{ - schema.AssistantMessage("p", nil), - }, streamErr), nil - }) - m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m2Calls, 1) - return streamWithMidError([]*schema.Message{ - schema.AssistantMessage("p", nil), - }, streamErr), nil - }) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 1, - IsRetryAble: func(_ context.Context, err error) bool { return true }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.Error(t, err) - - var retryErr *RetryExhaustedError - require.True(t, errors.As(err, &retryErr)) - - // m1: 1 initial + 1 retry = 2 calls - require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) - // m2: 1 initial + 1 retry = 2 calls - require.Equal(t, int32(2), atomic.LoadInt32(&m2Calls)) - }) - - t.Run("ShouldRetry_Stream_TriggersFailover", func(t *testing.T) { - var m1Calls int32 - var m2Calls int32 - - m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m1Calls, 1) - return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("bad from m1", nil)}), nil - }) - m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m2Calls, 1) - return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("good from m2", nil)}), nil - }) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 1, - ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { - if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad from m1" { - return &RetryDecision{Retry: true} - } - return &RetryDecision{Retry: false} - }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - msgs, err := drainMessageStream(sr) - require.NoError(t, err) - require.Len(t, msgs, 1) - require.Equal(t, "good from m2", msgs[0].Content) - require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) - }) - - t.Run("ShouldRetry_Generate_TriggersFailover", func(t *testing.T) { - var m1Calls int32 - var m2Calls int32 - - m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m1Calls, 1) - return schema.AssistantMessage("bad from m1", nil), nil - }, nil) - m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m2Calls, 1) - return schema.AssistantMessage("good from m2", nil), nil - }, nil) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 1, - ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { - if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad from m1" { - return &RetryDecision{Retry: true} - } - return &RetryDecision{Retry: false} - }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - return m2, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.NoError(t, err) - require.Equal(t, "good from m2", msg.Content) - require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) - }) - - t.Run("Stream_GetFailoverModelReturnsNilModel", func(t *testing.T) { - streamErr := errors.New("m1 always fails") - var m1Calls int32 - - m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m1Calls, 1) - return nil, streamErr - }) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 0, - IsRetryAble: func(_ context.Context, err error) bool { return false }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 1, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - return err != nil - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - return nil, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.Error(t, err) - require.Contains(t, err.Error(), "returned nil model at attempt") - require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) - }) - - t.Run("Stream_ContextCanceledDuringFailover", func(t *testing.T) { - streamErr := errors.New("m1 fails") - var m1Calls int32 - var failoverModelCalled int32 - - m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m1Calls, 1) - return nil, streamErr - }) - - ctx, cancel := context.WithCancel(context.Background()) - - retryCfg := &ModelRetryConfig{ - MaxRetries: 0, - IsRetryAble: func(_ context.Context, err error) bool { return false }, - BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, - } - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 3, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - cancel() - return err != nil - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - atomic.AddInt32(&failoverModelCalled, 1) - return nil, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - retryConfig: retryCfg, - failoverConfig: failoverCfg, - }) - - ctx = withTypedChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.Error(t, err) - require.ErrorIs(t, err, context.Canceled) - require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(0), atomic.LoadInt32(&failoverModelCalled)) - }) -} - -func TestErrStreamCanceled_Failover(t *testing.T) { - t.Run("Stream_NeverFailedOver", func(t *testing.T) { - var m1Calls int32 - var failoverCalled int32 - - m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { - atomic.AddInt32(&m1Calls, 1) - return streamWithMidError([]*schema.Message{ - schema.AssistantMessage("partial", nil), - }, ErrStreamCanceled), nil - }) - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 2, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - atomic.AddInt32(&failoverCalled, 1) - return true - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - t.Fatal("GetFailoverModel should not be called for ErrStreamCanceled") - return nil, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.Error(t, err) - require.True(t, errors.Is(err, ErrStreamCanceled)) - require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) - }) - - t.Run("Generate_NeverFailedOver", func(t *testing.T) { - var m1Calls int32 - var failoverCalled int32 - - m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { - atomic.AddInt32(&m1Calls, 1) - return nil, ErrStreamCanceled - }, nil) - - failoverCfg := &ModelFailoverConfig[*schema.Message]{ - MaxRetries: 2, - ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { - atomic.AddInt32(&failoverCalled, 1) - return true - }, - GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { - t.Fatal("GetFailoverModel should not be called for ErrStreamCanceled") - return nil, nil, nil - }, - } - - wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ - failoverConfig: failoverCfg, - }) - - ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ - failoverLastSuccessModel: m1, - }) - _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) - require.Error(t, err) - require.True(t, errors.Is(err, ErrStreamCanceled)) - require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) - require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) - }) -} diff --git a/adk/wrappers_test.go b/adk/wrappers_test.go index 11dacbd91..3e964115e 100644 --- a/adk/wrappers_test.go +++ b/adk/wrappers_test.go @@ -17,11 +17,15 @@ package adk import ( + "bytes" "context" + "encoding/gob" "errors" + "fmt" "sync" "sync/atomic" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -2002,9 +2006,8 @@ func TestTypedToolStreamEventAgenticMessageSetsStreamingMeta(t *testing.T) { require.Len(t, result.ContentBlocks, 1) assert.Nil(t, result.ContentBlocks[0].StreamingMeta) require.NotNil(t, result.ContentBlocks[0].FunctionToolResult) - require.Len(t, result.ContentBlocks[0].FunctionToolResult.Content, 2) - assert.Equal(t, "first\n", result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) - assert.Equal(t, "second\n", result.ContentBlocks[0].FunctionToolResult.Content[1].Text.Text) + require.Len(t, result.ContentBlocks[0].FunctionToolResult.Content, 1) + assert.Equal(t, "first\nsecond\n", result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) } func TestTypedToolEnhancedStreamEventAgenticMessageSetsStreamingMeta(t *testing.T) { @@ -2038,9 +2041,8 @@ func TestTypedToolEnhancedStreamEventAgenticMessageSetsStreamingMeta(t *testing. require.Len(t, result.ContentBlocks, 1) assert.Nil(t, result.ContentBlocks[0].StreamingMeta) require.NotNil(t, result.ContentBlocks[0].FunctionToolResult) - require.Len(t, result.ContentBlocks[0].FunctionToolResult.Content, 2) - assert.Equal(t, "first\n", result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) - assert.Equal(t, "second\n", result.ContentBlocks[0].FunctionToolResult.Content[1].Text.Text) + require.Len(t, result.ContentBlocks[0].FunctionToolResult.Content, 1) + assert.Equal(t, "first\nsecond\n", result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) } // multimodalEnhancedInvokableTestTool returns a pre-built multimodal ToolResult. @@ -2186,3 +2188,1411 @@ func TestExtractToolIdentifiersToolSearchResult(t *testing.T) { assert.Equal(t, "tool_search", toolName) assert.Equal(t, "call_1", callID) } + +func TestBuildModelWrappers_FailoverProxyInner(t *testing.T) { + base := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("ok", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil + }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 0, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return false }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return base, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](base, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + smw, ok := wrapped.(*stateModelWrapper) + require.True(t, ok) + _, ok = smw.inner.(*failoverProxyModel) + require.True(t, ok) + require.Same(t, base, smw.original) + require.Same(t, failoverCfg, smw.modelFailoverConfig) +} + +func TestStateModelWrapper_Generate_WithFailover(t *testing.T) { + wantErr := errors.New("first failed") + var shouldCalls int32 + var m1Calls int32 + var m2Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return schema.AssistantMessage("partial", nil), wantErr + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("ok", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + require.ErrorIs(t, err, wantErr) + require.NotNil(t, out) + require.Equal(t, "partial", out.Content) + return true + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + got, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.NotNil(t, got) + require.Equal(t, "ok", got.Content) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) +} + +func TestStateModelWrapper_Stream_WithFailover(t *testing.T) { + streamErr := errors.New("mid error") + var shouldCalls int32 + var m1Calls int32 + var m2Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("p1", nil), + schema.AssistantMessage("p2", nil), + }, streamErr), nil + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final", nil)}), nil + }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + require.ErrorIs(t, err, streamErr) + require.NotNil(t, out) + require.Equal(t, "p1p2", out.Content) + return true + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "final", msgs[0].Content) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) +} + +func TestFailoverAcceptsAgenticAgent(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("ok"), nil + }, + } + + fallbackModel := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("fallback"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "FailoverAgent", + Description: "Agent with failover config", + Model: m, + ModelFailoverConfig: &ModelFailoverConfig[*schema.AgenticMessage]{ + MaxRetries: 1, + ShouldFailover: func(ctx context.Context, outputMessage *schema.AgenticMessage, outputErr error) bool { + return true + }, + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext[*schema.AgenticMessage]) (model.BaseModel[*schema.AgenticMessage], []*schema.AgenticMessage, error) { + return fallbackModel, nil, nil + }, + }, + }) + require.NoError(t, err) + assert.NotNil(t, agent) +} + +// approvalInfoSpan and approvalResultSpan are isolated copies for use in this +// test file so we don't conflict with the prebuilt/integration_test.go types +// (which live in a different package anyway). +type approvalInfoSpan struct { + ToolName string + ArgumentsInJSON string + ToolCallID string +} + +type approvalResultSpan struct { + Approved bool +} + +func init() { + schema.Register[*approvalInfoSpan]() + schema.Register[*approvalResultSpan]() +} + +// approvableSpanTool is an invokable tool that interrupts on first invocation +// and runs to completion on resume after approval. +type approvableSpanTool struct { + name string +} + +func (t *approvableSpanTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "approvable span tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: schema.String, Desc: "input"}, + }), + }, nil +} + +func (t *approvableSpanTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + wasInterrupted, _, savedArgs := tool.GetInterruptState[string](ctx) + if !wasInterrupted { + return "", tool.StatefulInterrupt(ctx, &approvalInfoSpan{ + ToolName: t.name, + ArgumentsInJSON: argumentsInJSON, + ToolCallID: compose.GetToolCallID(ctx), + }, argumentsInJSON) + } + isResumeTarget, hasData, data := tool.GetResumeContext[*approvalResultSpan](ctx) + if !isResumeTarget || !hasData { + return "", tool.StatefulInterrupt(ctx, &approvalInfoSpan{ + ToolName: t.name, + ArgumentsInJSON: savedArgs, + ToolCallID: compose.GetToolCallID(ctx), + }, savedArgs) + } + if data.Approved { + return fmt.Sprintf("Tool '%s' executed with args: %s", t.name, savedArgs), nil + } + return fmt.Sprintf("Tool '%s' rejected", t.name), nil +} + +// approvableStreamableSpanTool: streamable variant. Interrupts on first +// invocation by returning a *core.InterruptSignal error before any stream +// chunk is produced; runs to completion on resume. +type approvableStreamableSpanTool struct { + name string +} + +func (t *approvableStreamableSpanTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "approvable streamable span tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: schema.String, Desc: "input"}, + }), + }, nil +} + +func (t *approvableStreamableSpanTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) { + wasInterrupted, _, savedArgs := tool.GetInterruptState[string](ctx) + if !wasInterrupted { + return nil, tool.StatefulInterrupt(ctx, &approvalInfoSpan{ + ToolName: t.name, + ArgumentsInJSON: argumentsInJSON, + ToolCallID: compose.GetToolCallID(ctx), + }, argumentsInJSON) + } + isResumeTarget, hasData, data := tool.GetResumeContext[*approvalResultSpan](ctx) + if !isResumeTarget || !hasData { + return nil, tool.StatefulInterrupt(ctx, &approvalInfoSpan{ + ToolName: t.name, + ArgumentsInJSON: savedArgs, + ToolCallID: compose.GetToolCallID(ctx), + }, savedArgs) + } + if data.Approved { + return schema.StreamReaderFromArray([]string{ + fmt.Sprintf("Tool '%s' streamed with args: %s", t.name, savedArgs), + }), nil + } + return schema.StreamReaderFromArray([]string{"rejected"}), nil +} + +// alwaysErrorTool errors out hard (non-interrupt) on every invocation. +type alwaysErrorTool struct { + name string +} + +func (t *alwaysErrorTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "always errors", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: schema.String, Desc: "input"}, + }), + }, nil +} + +func (t *alwaysErrorTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return "", errors.New("hard tool failure") +} + +// memCheckpointStore is a minimal in-memory CheckPointStore for these tests. +type memCheckpointStore struct { + mu sync.Mutex + data map[string][]byte +} + +func newMemCheckpointStore() *memCheckpointStore { + return &memCheckpointStore{data: make(map[string][]byte)} +} + +func (s *memCheckpointStore) Set(_ context.Context, key string, value []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + s.data[key] = value + return nil +} + +func (s *memCheckpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.data[key] + return v, ok, nil +} + +// scriptedToolCallingModel is a controllable mock model: each call returns the +// next scripted message. +type scriptedToolCallingModel struct { + mu sync.Mutex + messages []*schema.Message + pos int +} + +func (m *scriptedToolCallingModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.mu.Lock() + defer m.mu.Unlock() + if m.pos >= len(m.messages) { + return schema.AssistantMessage("done", nil), nil + } + msg := m.messages[m.pos] + m.pos++ + return msg, nil +} + +func (m *scriptedToolCallingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *scriptedToolCallingModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +// drainAndCollectSpans collects tool span events emitted during iter draining. +func drainAndCollectSpans(t *testing.T, iter *AsyncIterator[*AgentEvent]) (starts, ends []*SessionEvent[*schema.Message], interrupted bool) { + t.Helper() + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.Interrupted != nil { + interrupted = true + } + if ev.SessionEventVariant == nil || ev.SessionEventVariant.Event == nil || ev.SessionEventVariant.Event.Span == nil { + continue + } + switch ev.SessionEventVariant.Event.Kind { + case SessionEventSpanToolCallStart: + starts = append(starts, ev.SessionEventVariant.Event) + case SessionEventSpanToolCallEnd: + ends = append(ends, ev.SessionEventVariant.Event) + } + } + return +} + +// setupApprovableSpanAgent constructs a ChatModelAgent with a scripted +// tool-calling model and an in-memory checkpoint store, ready for span tests +// that exercise interrupt/resume. +func setupApprovableSpanAgent(t *testing.T, name string, tools []tool.BaseTool, scriptedAssistant []*schema.Message) (*TypedChatModelAgent[*schema.Message], *memCheckpointStore) { + t.Helper() + ctx := context.Background() + mdl := &scriptedToolCallingModel{messages: scriptedAssistant} + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: name, + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: tools}, + }, + }) + require.NoError(t, err) + store := newMemCheckpointStore() + return agent, store +} + +func TestToolSpan_PermissionInterruptDefersEndSpan(t *testing.T) { + ctx := context.Background() + tl := &approvableSpanTool{name: "approve_me"} + + scripted := []*schema.Message{ + schema.AssistantMessage("calling", []schema.ToolCall{{ID: "call_1", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"x"}`}}}), + } + agent, store := setupApprovableSpanAgent(t, "agent1", []tool.BaseTool{tl}, scripted) + runner := NewRunner(ctx, RunnerConfig{Agent: agent, CheckPointStore: store}) + + checkpointID := "ckpt-1" + iter := runner.Run(ctx, []Message{schema.UserMessage("go")}, WithCheckPointID(checkpointID), WithTimelineEvents()) + starts, ends, interrupted := drainAndCollectSpans(t, iter) + + assert.True(t, interrupted, "expected an interrupt event from the approvable tool") + require.Len(t, starts, 1, "exactly one tool_call_start span should be emitted on the interrupted run") + assert.Empty(t, ends, "no tool_call_end span should be emitted on the interrupted run") + assert.Equal(t, "call_1", starts[0].Span.Tool.ToolUseID) + assert.NotEmpty(t, starts[0].Span.Tool.AssistantMessageEventID, "start span must carry assistant message event ID") + assert.NotEmpty(t, starts[0].Span.ParentSpanID, "start span must carry parent (model) span ID") +} + +// runInterruptResumeAndCollectSpans drives a single-tool interrupt+approve +// scenario and returns the start span emitted on the original run plus the +// end span emitted on the resumed run. Used by both the resume happy-path +// test and the dedicated "parent IDs survive resume" assertion. +func runInterruptResumeAndCollectSpans(t *testing.T, agent *TypedChatModelAgent[*schema.Message], store *memCheckpointStore, checkpointID string) (startSpan, endSpan *SessionEvent[*schema.Message]) { + t.Helper() + ctx := context.Background() + runner := NewRunner(ctx, RunnerConfig{Agent: agent, CheckPointStore: store}) + iter1 := runner.Run(ctx, []Message{schema.UserMessage("go")}, WithCheckPointID(checkpointID), WithTimelineEvents()) + + var ( + starts1 []*SessionEvent[*schema.Message] + ends1 []*SessionEvent[*schema.Message] + interruptEvt *AgentEvent + ) + for { + ev, ok := iter1.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.Interrupted != nil { + interruptEvt = ev + } + if ev.SessionEventVariant != nil && ev.SessionEventVariant.Event != nil && ev.SessionEventVariant.Event.Span != nil { + switch ev.SessionEventVariant.Event.Kind { + case SessionEventSpanToolCallStart: + starts1 = append(starts1, ev.SessionEventVariant.Event) + case SessionEventSpanToolCallEnd: + ends1 = append(ends1, ev.SessionEventVariant.Event) + } + } + } + require.NotNil(t, interruptEvt) + require.Len(t, starts1, 1) + require.Empty(t, ends1) + + var toolInterruptID string + for _, ictx := range interruptEvt.Action.Interrupted.InterruptContexts { + if ictx.IsRootCause { + toolInterruptID = ictx.ID + break + } + } + require.NotEmpty(t, toolInterruptID) + + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &ResumeParams{ + Targets: map[string]any{toolInterruptID: &approvalResultSpan{Approved: true}}, + }, WithTimelineEvents()) + require.NoError(t, err) + _, ends2, _ := drainAndCollectSpans(t, resumeIter) + require.Len(t, ends2, 1) + return starts1[0], ends2[0] +} + +func TestToolSpan_PermissionResumeEmitsEndSpan(t *testing.T) { + tl := &approvableSpanTool{name: "approve_me"} + scripted := []*schema.Message{ + schema.AssistantMessage("calling", []schema.ToolCall{{ID: "call_resume", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"x"}`}}}), + schema.AssistantMessage("done", nil), + } + agent, store := setupApprovableSpanAgent(t, "agent_resume", []tool.BaseTool{tl}, scripted) + startSpan, endSpan := runInterruptResumeAndCollectSpans(t, agent, store, "ckpt-resume") + + assert.Equal(t, startSpan.Span.SpanID, endSpan.Span.SpanID, "end span must reuse the original SpanID") + assert.Equal(t, startSpan.EventID, endSpan.Span.Tool.ToolCallStartEventID, "end's ToolCallStartEventID must match start's EventID") + assert.Equal(t, "ok", endSpan.Span.Status) + assert.NotEmpty(t, endSpan.Span.Tool.ToolResultMessageEventID) +} + +// TestToolSpan_ResumeUsesOriginalTurnParentIDs (plan §4.5.1 #3) verifies that +// the resumed end span's ParentSpanID and AssistantMessageEventID match the +// original turn's model span and assistant message — confirming the in-flight +// span snapshot survived the checkpoint round-trip. +func TestToolSpan_ResumeUsesOriginalTurnParentIDs(t *testing.T) { + tl := &approvableSpanTool{name: "approve_me"} + scripted := []*schema.Message{ + schema.AssistantMessage("calling", []schema.ToolCall{{ID: "call_parents", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"x"}`}}}), + schema.AssistantMessage("done", nil), + } + agent, store := setupApprovableSpanAgent(t, "agent_parents", []tool.BaseTool{tl}, scripted) + startSpan, endSpan := runInterruptResumeAndCollectSpans(t, agent, store, "ckpt-parents") + + require.NotEmpty(t, startSpan.Span.ParentSpanID, "start span carries a non-empty parent (model) span ID") + require.NotEmpty(t, startSpan.Span.Tool.AssistantMessageEventID, "start span carries a non-empty assistant message event ID") + assert.Equal(t, startSpan.Span.ParentSpanID, endSpan.Span.ParentSpanID, "ParentSpanID survives resume via the in-flight snapshot") + assert.Equal(t, startSpan.Span.Tool.AssistantMessageEventID, endSpan.Span.Tool.AssistantMessageEventID, "AssistantMessageEventID survives resume via the in-flight snapshot") +} + +func TestToolSpan_HardErrorOnFirstRunStillEmitsEnd(t *testing.T) { + ctx := context.Background() + tl := &alwaysErrorTool{name: "boom"} + scripted := []*schema.Message{ + schema.AssistantMessage("calling", []schema.ToolCall{{ID: "err_call", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"x"}`}}}), + } + agent, store := setupApprovableSpanAgent(t, "err_agent", []tool.BaseTool{tl}, scripted) + runner := NewRunner(ctx, RunnerConfig{Agent: agent, CheckPointStore: store}) + iter := runner.Run(ctx, []Message{schema.UserMessage("go")}, WithCheckPointID("ckpt-err"), WithTimelineEvents()) + starts, ends, _ := drainAndCollectSpans(t, iter) + + require.Len(t, starts, 1) + require.Len(t, ends, 1) + assert.Equal(t, starts[0].Span.SpanID, ends[0].Span.SpanID, "end span shares SpanID with start span") + assert.Equal(t, "error", ends[0].Span.Status) +} + +func TestToolSpan_StreamableInterruptDefersEnd(t *testing.T) { + ctx := context.Background() + tl := &approvableStreamableSpanTool{name: "stream_approve_me"} + scripted := []*schema.Message{ + schema.AssistantMessage("calling", []schema.ToolCall{{ID: "stream_call", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"x"}`}}}), + schema.AssistantMessage("done", nil), + } + agent, store := setupApprovableSpanAgent(t, "stream_agent", []tool.BaseTool{tl}, scripted) + runner1 := NewRunner(ctx, RunnerConfig{Agent: agent, CheckPointStore: store}) + checkpointID := "ckpt-stream" + iter1 := runner1.Run(ctx, []Message{schema.UserMessage("go")}, WithCheckPointID(checkpointID), WithTimelineEvents()) + var interruptEvt *AgentEvent + starts1, ends1 := []*SessionEvent[*schema.Message]{}, []*SessionEvent[*schema.Message]{} + for { + ev, ok := iter1.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.Interrupted != nil { + interruptEvt = ev + } + if ev.SessionEventVariant != nil && ev.SessionEventVariant.Event != nil && ev.SessionEventVariant.Event.Span != nil { + switch ev.SessionEventVariant.Event.Kind { + case SessionEventSpanToolCallStart: + starts1 = append(starts1, ev.SessionEventVariant.Event) + case SessionEventSpanToolCallEnd: + ends1 = append(ends1, ev.SessionEventVariant.Event) + } + } + } + require.NotNil(t, interruptEvt) + require.Len(t, starts1, 1, "one start span on interrupted streamable run") + assert.Empty(t, ends1, "no end span on interrupted streamable run") + startSpanID := starts1[0].Span.SpanID + + var toolInterruptID string + for _, ictx := range interruptEvt.Action.Interrupted.InterruptContexts { + if ictx.IsRootCause { + toolInterruptID = ictx.ID + break + } + } + require.NotEmpty(t, toolInterruptID) + + resumeIter, err := runner1.ResumeWithParams(ctx, checkpointID, &ResumeParams{ + Targets: map[string]any{toolInterruptID: &approvalResultSpan{Approved: true}}, + }, WithTimelineEvents()) + require.NoError(t, err) + starts2, ends2, _ := drainAndCollectSpans(t, resumeIter) + assert.Empty(t, starts2, "no new start span on streamable resume") + require.Len(t, ends2, 1, "one end span on streamable resume") + assert.Equal(t, startSpanID, ends2[0].Span.SpanID, "end span shares SpanID with start span across resume") + assert.Equal(t, "ok", ends2[0].Span.Status) +} + +func TestTypedState_ToolSpansInFlightGobRoundTrip(t *testing.T) { + original := &typedState[*schema.Message]{ + Messages: []*schema.Message{schema.UserMessage("hello")}, + CurrentModelSpanID: "model-span-1", + CurrentAssistantMessageEventID: "asst-event-1", + ToolSpansInFlight: map[string]*toolSpanInFlight{ + "call_a": { + SpanID: "span-a", + StartEventID: "start-event-a", + StartedAt: time.Date(2026, 5, 26, 12, 0, 0, 0, time.UTC), + ParentSpanID: "model-span-1", + AssistantMessageEventID: "asst-event-1", + }, + "call_b": { + SpanID: "span-b", + StartEventID: "start-event-b", + StartedAt: time.Date(2026, 5, 26, 12, 0, 1, 0, time.UTC), + ParentSpanID: "model-span-1", + AssistantMessageEventID: "asst-event-1", + }, + }, + } + + var buf bytes.Buffer + require.NoError(t, gob.NewEncoder(&buf).Encode(original)) + + decoded := &typedState[*schema.Message]{} + require.NoError(t, gob.NewDecoder(&buf).Decode(decoded)) + + assert.Equal(t, original.CurrentModelSpanID, decoded.CurrentModelSpanID) + assert.Equal(t, original.CurrentAssistantMessageEventID, decoded.CurrentAssistantMessageEventID) + require.Len(t, decoded.ToolSpansInFlight, 2) + for k, v := range original.ToolSpansInFlight { + got, ok := decoded.ToolSpansInFlight[k] + require.Truef(t, ok, "missing key %q after gob round-trip", k) + assert.Equal(t, v.SpanID, got.SpanID) + assert.Equal(t, v.StartEventID, got.StartEventID) + assert.True(t, v.StartedAt.Equal(got.StartedAt), "StartedAt mismatch: %v vs %v", v.StartedAt, got.StartedAt) + assert.Equal(t, v.ParentSpanID, got.ParentSpanID) + assert.Equal(t, v.AssistantMessageEventID, got.AssistantMessageEventID) + } +} + +// Sanity guard: ensure compose.IsInterruptRerunError import is preserved (used in wrappers). +var _ = compose.IsInterruptRerunError + +// TestToolSpan_PermissionRejectEmitsEndSpan exercises the path where the tool +// is interrupted, then on resume the user rejects (Approved=false). The tool +// returns a rejection result rather than an error, so the end span carries +// Status=ok with a populated ToolResultMessageEventID. Same SpanID across +// the boundary. +func TestToolSpan_PermissionRejectEmitsEndSpan(t *testing.T) { + ctx := context.Background() + tl := &approvableSpanTool{name: "reject_me"} + + scripted := []*schema.Message{ + schema.AssistantMessage("calling", []schema.ToolCall{{ID: "rej_call", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"x"}`}}}), + schema.AssistantMessage("done", nil), + } + agent, store := setupApprovableSpanAgent(t, "reject_agent", []tool.BaseTool{tl}, scripted) + checkpointID := "ckpt-reject" + runner := NewRunner(ctx, RunnerConfig{Agent: agent, CheckPointStore: store}) + iter1 := runner.Run(ctx, []Message{schema.UserMessage("go")}, WithCheckPointID(checkpointID), WithTimelineEvents()) + + var ( + starts1 []*SessionEvent[*schema.Message] + ends1 []*SessionEvent[*schema.Message] + interruptEvt *AgentEvent + ) + for { + ev, ok := iter1.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.Interrupted != nil { + interruptEvt = ev + } + if ev.SessionEventVariant != nil && ev.SessionEventVariant.Event != nil && ev.SessionEventVariant.Event.Span != nil { + switch ev.SessionEventVariant.Event.Kind { + case SessionEventSpanToolCallStart: + starts1 = append(starts1, ev.SessionEventVariant.Event) + case SessionEventSpanToolCallEnd: + ends1 = append(ends1, ev.SessionEventVariant.Event) + } + } + } + require.NotNil(t, interruptEvt) + require.Len(t, starts1, 1) + require.Empty(t, ends1) + + startSpanID := starts1[0].Span.SpanID + + var toolInterruptID string + for _, ictx := range interruptEvt.Action.Interrupted.InterruptContexts { + if ictx.IsRootCause { + toolInterruptID = ictx.ID + break + } + } + require.NotEmpty(t, toolInterruptID) + + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &ResumeParams{ + Targets: map[string]any{toolInterruptID: &approvalResultSpan{Approved: false}}, + }, WithTimelineEvents()) + require.NoError(t, err) + starts2, ends2, _ := drainAndCollectSpans(t, resumeIter) + assert.Empty(t, starts2) + require.Len(t, ends2, 1) + assert.Equal(t, startSpanID, ends2[0].Span.SpanID) + assert.Equal(t, "ok", ends2[0].Span.Status, "rejection produces a successful return (the deny content) — status is ok, not error") + assert.NotEmpty(t, ends2[0].Span.Tool.ToolResultMessageEventID) +} + +// successOnlyTool runs to completion on the first invocation. Combined with +// the absence of a permission middleware, it exercises the "non-interrupted +// call" path where start and end both fire on the same run with status=ok. +type successOnlyTool struct { + name string + result string +} + +func (t *successOnlyTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "always succeeds", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: schema.String, Desc: "input"}, + }), + }, nil +} + +func (t *successOnlyTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return t.result, nil +} + +// TestToolSpan_NonInterruptedCallEmitsBothSpansOnSameRun verifies that a tool +// that runs straight to success produces a tool_call_start + tool_call_end +// pair on the same run, with the in-flight entry cleared at end emission. +// (This corresponds to plan §4.5.1 #6 which used a "gate=deny" example — +// the wire-shape behavior is identical: single run, single span pair, status +// ok, populated ToolResultMessageEventID.) +func TestToolSpan_NonInterruptedCallEmitsBothSpansOnSameRun(t *testing.T) { + ctx := context.Background() + tl := &successOnlyTool{name: "noninterrupted_tool", result: "ok"} + scripted := []*schema.Message{ + schema.AssistantMessage("calling", []schema.ToolCall{{ID: "noninter_call", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"x"}`}}}), + schema.AssistantMessage("done", nil), + } + agent, store := setupApprovableSpanAgent(t, "noninter_agent", []tool.BaseTool{tl}, scripted) + runner := NewRunner(ctx, RunnerConfig{Agent: agent, CheckPointStore: store}) + iter := runner.Run(ctx, []Message{schema.UserMessage("go")}, WithCheckPointID("ckpt-noninter"), WithTimelineEvents()) + starts, ends, _ := drainAndCollectSpans(t, iter) + + require.Len(t, starts, 1) + require.Len(t, ends, 1) + assert.Equal(t, starts[0].Span.SpanID, ends[0].Span.SpanID) + assert.Equal(t, "ok", ends[0].Span.Status) + assert.NotEmpty(t, ends[0].Span.Tool.ToolResultMessageEventID) +} + +// TestToolSpan_ParallelInterruptResumesEmitMatchingEnds exercises the parallel +// call scenario: two tool calls (A, B) emitted in a single assistant message, +// both interrupting on first invocation. After the first run we should see +// 2 starts and 0 ends. After resuming both with approval, we expect end spans +// keyed to the matching SpanIDs (one per CallID). +func TestToolSpan_ParallelInterruptResumesEmitMatchingEnds(t *testing.T) { + ctx := context.Background() + tl := &approvableSpanTool{name: "parallel_tool"} + scripted := []*schema.Message{ + schema.AssistantMessage("calling 2", []schema.ToolCall{ + {ID: "call_par_a", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"a"}`}}, + {ID: "call_par_b", Function: schema.FunctionCall{Name: tl.name, Arguments: `{"input":"b"}`}}, + }), + schema.AssistantMessage("done", nil), + } + agent, store := setupApprovableSpanAgent(t, "parallel_agent", []tool.BaseTool{tl}, scripted) + checkpointID := "ckpt-parallel" + runner := NewRunner(ctx, RunnerConfig{Agent: agent, CheckPointStore: store}) + iter1 := runner.Run(ctx, []Message{schema.UserMessage("go")}, WithCheckPointID(checkpointID), WithTimelineEvents()) + + var ( + starts1 []*SessionEvent[*schema.Message] + ends1 []*SessionEvent[*schema.Message] + interruptEvt *AgentEvent + ) + for { + ev, ok := iter1.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.Interrupted != nil { + interruptEvt = ev + } + if ev.SessionEventVariant != nil && ev.SessionEventVariant.Event != nil && ev.SessionEventVariant.Event.Span != nil { + switch ev.SessionEventVariant.Event.Kind { + case SessionEventSpanToolCallStart: + starts1 = append(starts1, ev.SessionEventVariant.Event) + case SessionEventSpanToolCallEnd: + ends1 = append(ends1, ev.SessionEventVariant.Event) + } + } + } + require.NotNil(t, interruptEvt) + require.Len(t, starts1, 2, "expected one tool_call_start for each parallel call") + assert.Empty(t, ends1) + + // Map CallID -> start SpanID for later assertions. + callIDToStartSpanID := map[string]string{} + for _, s := range starts1 { + callIDToStartSpanID[s.Span.Tool.ToolUseID] = s.Span.SpanID + } + require.Contains(t, callIDToStartSpanID, "call_par_a") + require.Contains(t, callIDToStartSpanID, "call_par_b") + + // Collect interrupt IDs (root causes only). + var interruptIDs []string + for _, ictx := range interruptEvt.Action.Interrupted.InterruptContexts { + if ictx.IsRootCause { + interruptIDs = append(interruptIDs, ictx.ID) + } + } + require.Len(t, interruptIDs, 2) + + // Approve both at once. + targets := map[string]any{} + for _, id := range interruptIDs { + targets[id] = &approvalResultSpan{Approved: true} + } + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &ResumeParams{Targets: targets}, WithTimelineEvents()) + require.NoError(t, err) + starts2, ends2, _ := drainAndCollectSpans(t, resumeIter) + assert.Empty(t, starts2, "no new starts on resume") + require.Len(t, ends2, 2, "two ends, one per parallel call") + for _, e := range ends2 { + expectedSpanID, ok := callIDToStartSpanID[e.Span.Tool.ToolUseID] + require.Truef(t, ok, "end span carries unknown CallID %q", e.Span.Tool.ToolUseID) + assert.Equal(t, expectedSpanID, e.Span.SpanID, "end span SpanID matches the start span for the same CallID") + assert.Equal(t, "ok", e.Span.Status) + } +} + +func newFakeChatModel( + gen func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error), + stream func(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error), +) *fakeChatModel { + if gen == nil { + gen = func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + } + } + if stream == nil { + stream = func(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + } + } + return &fakeChatModel{callbacksEnabled: true, generate: gen, stream: stream} +} + +func TestRetryThenFailover(t *testing.T) { + t.Run("Generate_RetryExhaustedTriggersFailover", func(t *testing.T) { + modelErr := errors.New("model error") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, modelErr + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("ok from m2", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 2, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, fc *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.NotNil(t, fc.LastErr) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "ok from m2", msg.Content) + + // m1: 1 (lastSuccess) + 2 retries = 3 calls on lastSuccess attempt, + // then failover to m2 which also goes through retry wrapper: 1 call succeeds. + require.Equal(t, int32(3), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Generate_AllExhausted", func(t *testing.T) { + modelErr := errors.New("always fails") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, modelErr + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return nil, modelErr + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + + // Should be RetryExhaustedError from m2's retry wrapper + var retryErr *RetryExhaustedError + require.True(t, errors.As(err, &retryErr)) + + // m1: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + // m2: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Generate_RetrySucceedsNoFailover", func(t *testing.T) { + var m1Calls int32 + var failoverCalled int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + n := atomic.AddInt32(&m1Calls, 1) + if n == 1 { + return nil, errors.New("transient error") + } + return schema.AssistantMessage("ok on retry", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 2, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&failoverCalled, 1) + return true + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + t.Fatal("GetFailoverModel should not be called when retry succeeds") + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "ok on retry", msg.Content) + + // 2 calls: first fails, second succeeds via retry + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + // ShouldFailover should never be called + require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) + }) + + t.Run("Generate_NonRetryableErrorTriggersFailover", func(t *testing.T) { + nonRetryableErr := errors.New("non-retryable") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, nonRetryableErr + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("ok from m2", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(_ context.Context, err error) bool { + // Only non-retryable errors + return !errors.Is(err, nonRetryableErr) + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "ok from m2", msg.Content) + + // m1 called only once — non-retryable error skips retry + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Stream_RetryExhaustedTriggersFailover", func(t *testing.T) { + streamErr := errors.New("stream mid error") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("partial", nil), + }, streamErr), nil + }) + m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok from m2", nil)}), nil + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, fc *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.NotNil(t, fc.LastErr) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "ok from m2", msgs[0].Content) + + // m1: 1 initial + 1 retry = 2 calls on lastSuccess attempt + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Stream_AllExhausted", func(t *testing.T) { + streamErr := errors.New("always fails mid-stream") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("p", nil), + }, streamErr), nil + }) + m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("p", nil), + }, streamErr), nil + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + + var retryErr *RetryExhaustedError + require.True(t, errors.As(err, &retryErr)) + + // m1: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + // m2: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("ShouldRetry_Stream_TriggersFailover", func(t *testing.T) { + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("bad from m1", nil)}), nil + }) + m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("good from m2", nil)}), nil + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad from m1" { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "good from m2", msgs[0].Content) + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("ShouldRetry_Generate_TriggersFailover", func(t *testing.T) { + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return schema.AssistantMessage("bad from m1", nil), nil + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("good from m2", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad from m1" { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "good from m2", msg.Content) + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Stream_GetFailoverModelReturnsNilModel", func(t *testing.T) { + streamErr := errors.New("m1 always fails") + var m1Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return nil, streamErr + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 0, + IsRetryAble: func(_ context.Context, err error) bool { return false }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.Contains(t, err.Error(), "returned nil model at attempt") + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + }) + + t.Run("Stream_ContextCanceledDuringFailover", func(t *testing.T) { + streamErr := errors.New("m1 fails") + var m1Calls int32 + var failoverModelCalled int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return nil, streamErr + }) + + ctx, cancel := context.WithCancel(context.Background()) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 0, + IsRetryAble: func(_ context.Context, err error) bool { return false }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 3, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + cancel() + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + atomic.AddInt32(&failoverModelCalled, 1) + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx = withTypedChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&failoverModelCalled)) + }) +} + +func TestErrStreamCanceled_Failover(t *testing.T) { + t.Run("Stream_NeverFailedOver", func(t *testing.T) { + var m1Calls int32 + var failoverCalled int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("partial", nil), + }, ErrStreamCanceled), nil + }) + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 2, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&failoverCalled, 1) + return true + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + t.Fatal("GetFailoverModel should not be called for ErrStreamCanceled") + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.True(t, errors.Is(err, ErrStreamCanceled)) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) + }) + + t.Run("Generate_NeverFailedOver", func(t *testing.T) { + var m1Calls int32 + var failoverCalled int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, ErrStreamCanceled + }, nil) + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 2, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&failoverCalled, 1) + return true + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + t.Fatal("GetFailoverModel should not be called for ErrStreamCanceled") + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.True(t, errors.Is(err, ErrStreamCanceled)) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) + }) +} diff --git a/compose/agentic_tools_node_test.go b/compose/agentic_tools_node_test.go index 1f5796304..2947a62bc 100644 --- a/compose/agentic_tools_node_test.go +++ b/compose/agentic_tools_node_test.go @@ -17,6 +17,7 @@ package compose import ( + "context" "io" "testing" @@ -24,6 +25,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/schema" ) @@ -343,6 +345,57 @@ func TestStreamToolMessageToAgenticMessage(t *testing.T) { }) } +func TestAgenticToolsNodeStreamSetsStreamingMeta(t *testing.T) { + ctx := context.Background() + node, err := NewAgenticToolsNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{&mockTool{}}, + }) + require.NoError(t, err) + + stream, err := node.Stream(ctx, &schema.AgenticMessage{ + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: &schema.FunctionToolCall{ + CallID: "call_1", + Name: "mock_tool", + Arguments: `{"name":"jack"}`, + }, + }, + }, + }) + require.NoError(t, err) + defer stream.Close() + + var chunks [][]*schema.AgenticMessage + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + require.NoError(t, err) + require.Len(t, chunk, 1) + require.Len(t, chunk[0].ContentBlocks, 1) + block := chunk[0].ContentBlocks[0] + assert.Equal(t, schema.ContentBlockTypeFunctionToolResult, block.Type) + assert.Equal(t, &schema.StreamingMeta{Index: 0}, block.StreamingMeta) + chunks = append(chunks, chunk) + } + require.NotEmpty(t, chunks) + + result, err := schema.ConcatAgenticMessagesArray(chunks) + require.NoError(t, err) + require.Len(t, result, 1) + require.Len(t, result[0].ContentBlocks, 1) + block := result[0].ContentBlocks[0] + assert.Nil(t, block.StreamingMeta) + require.NotNil(t, block.FunctionToolResult) + assert.Equal(t, "call_1", block.FunctionToolResult.CallID) + assert.Equal(t, "mock_tool", block.FunctionToolResult.Name) + require.Len(t, block.FunctionToolResult.Content, 1) + assert.JSONEq(t, `{"echo":"jack: 0"}`, block.FunctionToolResult.Content[0].Text.Text) +} + func testStreamToolMessageTextOnly(t *testing.T) { input := schema.StreamReaderFromArray([][]*schema.Message{ { @@ -435,8 +488,7 @@ func testStreamToolMessageTextOnly(t *testing.T) { CallID: "2", Name: "name2", Content: []*schema.FunctionToolResultContentBlock{ - {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content2-1"}}, - {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content2-2"}}, + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content2-1content2-2"}}, }, }, }, @@ -451,8 +503,7 @@ func testStreamToolMessageTextOnly(t *testing.T) { CallID: "3", Name: "name3", Content: []*schema.FunctionToolResultContentBlock{ - {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content3-1"}}, - {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content3-2"}}, + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content3-1content3-2"}}, }, }, }, diff --git a/compose/checkpoint.go b/compose/checkpoint.go index c174994d2..820268dde 100644 --- a/compose/checkpoint.go +++ b/compose/checkpoint.go @@ -51,10 +51,7 @@ func RegisterSerializableType[T any](name string) (err error) { type CheckPointStore = core.CheckPointStore -type Serializer interface { - Marshal(v any) ([]byte, error) - Unmarshal(data []byte, v any) error -} +type Serializer = schema.Serializer // WithCheckPointStore sets the checkpoint store implementation for a graph. func WithCheckPointStore(store CheckPointStore) GraphCompileOption { diff --git a/compose/graph_manager.go b/compose/graph_manager.go index 46df3488e..7bf031734 100644 --- a/compose/graph_manager.go +++ b/compose/graph_manager.go @@ -255,15 +255,16 @@ func appendIfNotExist(s []string, elem string) []string { } type task struct { - ctx context.Context - nodeKey string - call *chanCall - input any - originalInput any - output any - option []any - err error - skipPreHandler bool + ctx context.Context + nodeKey string + call *chanCall + input any + originalInput any + output any + option []any + err error + skipPreHandler bool + syntheticRerunInput bool } type taskManager struct { @@ -308,7 +309,7 @@ func (t *taskManager) submit(tasks []*task) error { for i := 0; i < len(tasks); i++ { currentTask := tasks[i] - if t.persistRerunInput { + if t.persistRerunInput && !currentTask.syntheticRerunInput { if sr, ok := currentTask.input.(streamReader); ok { copies := sr.copy(2) currentTask.originalInput, currentTask.input = copies[0], copies[1] diff --git a/compose/graph_run.go b/compose/graph_run.go index 02b4fca7d..434f6c4ac 100644 --- a/compose/graph_run.go +++ b/compose/graph_run.go @@ -21,6 +21,8 @@ import ( "errors" "fmt" "reflect" + "runtime/debug" + "sort" "strings" "github.com/cloudwego/eino/internal" @@ -560,7 +562,7 @@ func (r *runner) handleInterrupt( } else if checkPointID != nil { err := r.checkPointer.set(ctx, *checkPointID, cp) if err != nil { - return fmt.Errorf("failed to set checkpoint: %w, checkPointID: %s", err, *checkPointID) + return newCheckpointSetError("interrupt", *checkPointID, cp, err) } } @@ -700,13 +702,138 @@ func (r *runner) handleInterruptWithSubGraphAndRerunNodes( } else if checkPointID != nil { err = r.checkPointer.set(ctx, *checkPointID, cp) if err != nil { - return fmt.Errorf("failed to set checkpoint: %w, checkPointID: %s", err, *checkPointID) + return newCheckpointSetError("interrupt_with_subgraph_and_rerun_nodes", *checkPointID, cp, err) } } intInfo.InterruptContexts = core.ToInterruptContexts(is, nil) return &interruptError{Info: intInfo} } +const checkpointDebugEntryLimit = 8 + +func newCheckpointSetError(stage string, checkPointID string, cp *checkpoint, err error) error { + return fmt.Errorf("failed to set checkpoint during %s: %w, checkPointID: %s, checkpoint: %s, stack:\n%s", + stage, err, checkPointID, checkpointDebugSummary(cp), string(debug.Stack())) +} + +func checkpointDebugSummary(cp *checkpoint) string { + if cp == nil { + return "" + } + + var b strings.Builder + appendCheckpointDebug(&b, "root", cp, 0) + return b.String() +} + +func appendCheckpointDebug(b *strings.Builder, path string, cp *checkpoint, depth int) { + if cp == nil { + fmt.Fprintf(b, "%s=", path) + return + } + + fmt.Fprintf(b, "%s{state=%s rerunNodes=%v skipPreHandler=%v interruptAddr=%d interruptState=%d ", + path, + valueDebugSummary(cp.State), + cp.RerunNodes, + cp.SkipPreHandler, + len(cp.InterruptID2Addr), + len(cp.InterruptID2State)) + appendAnyMapDebug(b, "inputs", cp.Inputs) + b.WriteByte(' ') + appendChannelsDebug(b, cp.Channels) + b.WriteByte(' ') + appendSubGraphsDebug(b, cp.SubGraphs, path, depth) + b.WriteByte('}') +} + +func appendAnyMapDebug(b *strings.Builder, label string, values map[string]any) { + fmt.Fprintf(b, "%s(count=%d", label, len(values)) + for _, key := range sortedKeys(values, checkpointDebugEntryLimit) { + fmt.Fprintf(b, " %s=%s", key, valueDebugSummary(values[key])) + } + appendOmittedCount(b, len(values)) + b.WriteByte(')') +} + +func appendChannelsDebug(b *strings.Builder, channels map[string]channel) { + fmt.Fprintf(b, "channels(count=%d", len(channels)) + for _, key := range sortedKeys(channels, checkpointDebugEntryLimit) { + ch := channels[key] + fmt.Fprintf(b, " %s=%T", key, ch) + if ch == nil { + continue + } + + err := ch.convertValues(func(values map[string]any) error { + b.WriteByte('[') + for _, valueKey := range sortedKeys(values, checkpointDebugEntryLimit) { + fmt.Fprintf(b, "%s=%s", valueKey, valueDebugSummary(values[valueKey])) + } + appendOmittedCount(b, len(values)) + b.WriteByte(']') + return nil + }) + if err != nil { + fmt.Fprintf(b, "[convertValuesErr=%v]", err) + } + } + appendOmittedCount(b, len(channels)) + b.WriteByte(')') +} + +func appendSubGraphsDebug(b *strings.Builder, subGraphs map[string]*checkpoint, path string, depth int) { + fmt.Fprintf(b, "subGraphs(count=%d", len(subGraphs)) + if depth >= 2 { + if len(subGraphs) > 0 { + b.WriteString(" ...") + } + b.WriteByte(')') + return + } + + for _, key := range sortedKeys(subGraphs, checkpointDebugEntryLimit) { + b.WriteByte(' ') + appendCheckpointDebug(b, path+"."+key, subGraphs[key], depth+1) + } + appendOmittedCount(b, len(subGraphs)) + b.WriteByte(')') +} + +func valueDebugSummary(v any) string { + if v == nil { + return "" + } + + rv := reflect.ValueOf(v) + nilLike := false + switch rv.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice: + nilLike = rv.IsNil() + } + + _, isStream := v.(streamReader) + return fmt.Sprintf("%T(nil=%t stream=%t)", v, nilLike, isStream) +} + +func sortedKeys[V any](m map[string]V, limit int) []string { + keys := make([]string, 0, len(m)) + for key := range m { + keys = append(keys, key) + } + sort.Strings(keys) + if len(keys) > limit { + return keys[:limit] + } + return keys +} + +func appendOmittedCount(b *strings.Builder, total int) { + if total > checkpointDebugEntryLimit { + fmt.Fprintf(b, " ...+%d", total-checkpointDebugEntryLimit) + } +} + func (r *runner) calculateNextTasks(ctx context.Context, completedTasks []*task, isStream bool, cm *channelManager, optMap map[string][]any) ([]*task, any, bool, error) { writeChannelValues, controls, err := r.resolveCompletedTasks(ctx, completedTasks, isStream, cm) if err != nil { @@ -782,6 +909,7 @@ func (r *runner) restoreTasks( isStream bool, optMap map[string][]any) ([]*task, error) { ret := make([]*task, 0, len(inputs)) + syntheticInputs := make(map[string]struct{}) for _, key := range rerunNodes { if _, hasInput := inputs[key]; hasInput { continue @@ -796,6 +924,7 @@ func (r *runner) restoreTasks( } else { inputs[key] = call.action.inputZeroValue() } + syntheticInputs[key] = struct{}{} } for key, input := range inputs { call, ok := r.chanSubscribeTo[key] @@ -816,6 +945,9 @@ func (r *runner) restoreTasks( option: nil, skipPreHandler: skipPreHandler[key], } + if _, ok := syntheticInputs[key]; ok { + newTask.syntheticRerunInput = true + } if opt, ok := optMap[key]; ok { newTask.option = opt } diff --git a/examples b/examples deleted file mode 160000 index a51a4a8e6..000000000 --- a/examples +++ /dev/null @@ -1 +0,0 @@ -Subproject commit a51a4a8e6d9982eebdbf60a6518bdbde7a07dd45 diff --git a/ext b/ext deleted file mode 160000 index 8c43b097e..000000000 --- a/ext +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 8c43b097ea865c91927d73417bf10c19ff25e680 diff --git a/internal/serialization/gob_serializer.go b/internal/serialization/gob_serializer.go new file mode 100644 index 000000000..80cee5ced --- /dev/null +++ b/internal/serialization/gob_serializer.go @@ -0,0 +1,48 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package serialization + +import ( + "bytes" + "encoding/gob" + "fmt" + "reflect" +) + +type GobSerializer struct{} + +func (g *GobSerializer) Marshal(v any) ([]byte, error) { + var buf bytes.Buffer + enc := gob.NewEncoder(&buf) + if err := enc.Encode(v); err != nil { + return nil, fmt.Errorf("gob marshal error: %w", err) + } + return buf.Bytes(), nil +} + +func (g *GobSerializer) Unmarshal(data []byte, v any) error { + rv := reflect.ValueOf(v) + if rv.Kind() != reflect.Ptr || rv.IsNil() { + return fmt.Errorf("unmarshal destination must be a non-nil pointer") + } + + dec := gob.NewDecoder(bytes.NewReader(data)) + if err := dec.Decode(v); err != nil { + return fmt.Errorf("gob unmarshal error: %w", err) + } + return nil +} diff --git a/internal/serialization/human_readable.go b/internal/serialization/human_readable.go new file mode 100644 index 000000000..a9ed81129 --- /dev/null +++ b/internal/serialization/human_readable.go @@ -0,0 +1,957 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package serialization + +import ( + "bytes" + "encoding/json" + "fmt" + "reflect" + "strconv" + "strings" + + "github.com/bytedance/sonic" +) + +const typeFieldName = "$type" + +type HumanReadableSerializer struct{} + +func (h *HumanReadableSerializer) Marshal(v any) ([]byte, error) { + result, err := hrMarshal(v, nil) + if err != nil { + return nil, err + } + return sonic.Marshal(result) +} + +func (h *HumanReadableSerializer) Unmarshal(data []byte, v any) error { + rv := reflect.ValueOf(v) + if rv.Kind() != reflect.Ptr || rv.IsNil() { + return fmt.Errorf("failed to unmarshal: value must be a non-nil pointer") + } + + raw, err := decodeJSONAny(data) + if err != nil { + return fmt.Errorf("failed to unmarshal JSON: %w", err) + } + + result, err := hrUnmarshal(raw, rv.Elem().Type()) + if err != nil { + return fmt.Errorf("failed to unmarshal: %w", err) + } + + target := rv.Elem() + if !target.CanSet() { + return fmt.Errorf("failed to unmarshal: output value must be settable") + } + + if result == nil { + target.Set(reflect.Zero(target.Type())) + return nil + } + + source := reflect.ValueOf(result) + if !setValueWithConversion(target, source) { + return fmt.Errorf("failed to unmarshal: cannot assign %s to %s", reflect.TypeOf(result), target.Type()) + } + + return nil +} + +func hrMarshal(v any, fieldType reflect.Type) (any, error) { + if v == nil { + return nil, nil + } + + rv := reflect.ValueOf(v) + if !rv.IsValid() { + return nil, nil + } + + if rv.Kind() == reflect.Invalid { + return nil, nil + } + + rt := rv.Type() + typeUnspecific := fieldType == nil || fieldType.Kind() == reflect.Interface + + var pointerNum uint32 + for rt.Kind() == reflect.Ptr { + pointerNum++ + if rv.IsNil() { + return nil, nil + } + rv = rv.Elem() + rt = rt.Elem() + } + + switch rt.Kind() { + case reflect.Struct: + return hrMarshalStruct(rv, rt, typeUnspecific, pointerNum) + case reflect.Map: + return hrMarshalMap(rv, rt, typeUnspecific, pointerNum) + case reflect.Slice, reflect.Array: + return hrMarshalSlice(rv, rt, typeUnspecific, pointerNum) + default: + return hrMarshalPrimitive(rv, rt, typeUnspecific, pointerNum) + } +} + +func hrMarshalStruct(rv reflect.Value, rt reflect.Type, typeUnspecific bool, pointerNum uint32) (any, error) { + if checkMarshaler(rt) { + // Use the addressable form when possible so pointer-receiver MarshalJSON + // methods are invoked. Without this, custom marshalers defined on *T are + // silently bypassed because rv.Interface() returns a non-addressable copy + // and the standard json package only checks Marshaler on the value type. + // Symptom: types like ToolInfo that store data behind unexported fields + // (ParamsOneOf.params / .jsonschema) round-trip with those fields lost. + var marshalTarget any + if rv.CanAddr() { + marshalTarget = rv.Addr().Interface() + } else { + tmp := reflect.New(rt) + tmp.Elem().Set(rv) + marshalTarget = tmp.Interface() + } + jsonBytes, err := json.Marshal(marshalTarget) + if err != nil { + return nil, err + } + result, err := decodeJSONAny(jsonBytes) + if err != nil { + return nil, err + } + if typeUnspecific { + _, isMap := result.(map[string]any) + return wrapWithType(result, rt, pointerNum, !isMap) + } + return result, nil + } + + result := make(map[string]any) + + for i := 0; i < rt.NumField(); i++ { + field := rt.Field(i) + if field.PkgPath != "" { + continue + } + + fieldValue := rv.Field(i) + jsonTag := field.Tag.Get("json") + fieldName := getJSONFieldName(field.Name, jsonTag) + + if fieldName == "-" { + continue + } + + if hasOmitempty(jsonTag) && isEmptyValue(fieldValue) { + continue + } + + marshaledValue, err := hrMarshal(fieldValue.Interface(), field.Type) + if err != nil { + return nil, fmt.Errorf("failed to marshal field %s: %w", field.Name, err) + } + + if marshaledValue != nil || !hasOmitempty(jsonTag) { + result[fieldName] = marshaledValue + } + } + + if typeUnspecific { + return wrapWithType(result, rt, pointerNum, false) + } + + return result, nil +} + +func hrMarshalMap(rv reflect.Value, rt reflect.Type, typeUnspecific bool, pointerNum uint32) (any, error) { + if rv.IsNil() { + return nil, nil + } + + result := make(map[string]any) + iter := rv.MapRange() + + for iter.Next() { + k := iter.Key() + v := iter.Value() + + var keyStr string + if k.Kind() == reflect.String { + keyStr = k.String() + } else { + keyBytes, err := sonic.Marshal(k.Interface()) + if err != nil { + return nil, fmt.Errorf("failed to marshal map key: %w", err) + } + keyStr = string(keyBytes) + } + + marshaledValue, err := hrMarshal(v.Interface(), rt.Elem()) + if err != nil { + return nil, fmt.Errorf("failed to marshal map value for key %s: %w", keyStr, err) + } + + result[keyStr] = marshaledValue + } + + if typeUnspecific { + return wrapMapWithType(result, rt, pointerNum) + } + + return result, nil +} + +func hrMarshalSlice(rv reflect.Value, rt reflect.Type, typeUnspecific bool, pointerNum uint32) (any, error) { + if rv.Kind() == reflect.Slice && rv.IsNil() { + return nil, nil + } + + length := rv.Len() + result := make([]any, length) + + for i := 0; i < length; i++ { + elem := rv.Index(i) + marshaledElem, err := hrMarshal(elem.Interface(), rt.Elem()) + if err != nil { + return nil, fmt.Errorf("failed to marshal slice element %d: %w", i, err) + } + result[i] = marshaledElem + } + + if typeUnspecific { + return wrapSliceWithType(result, rt, pointerNum) + } + + return result, nil +} + +func hrMarshalPrimitive(rv reflect.Value, rt reflect.Type, typeUnspecific bool, pointerNum uint32) (any, error) { + if !typeUnspecific { + return rv.Interface(), nil + } + + if isPrimitiveJSONType(rt) && pointerNum == 0 { + return rv.Interface(), nil + } + + return wrapWithType(rv.Interface(), rt, pointerNum, true) +} + +func wrapWithType(value any, rt reflect.Type, pointerNum uint32, isSimple bool) (any, error) { + key, ok := rm[rt] + if !ok { + return nil, fmt.Errorf("unknown type: %v (not registered)", rt) + } + + typeName := key + if pointerNum > 0 { + typeName = strings.Repeat("*", int(pointerNum)) + key + } + + if isSimple { + return map[string]any{ + typeFieldName: typeName, + "value": value, + }, nil + } + + if m, ok := value.(map[string]any); ok { + if _, exists := m[typeFieldName]; !exists { + m[typeFieldName] = typeName + return m, nil + } + } + + return map[string]any{ + typeFieldName: typeName, + "value": value, + }, nil +} + +func wrapMapWithType(value map[string]any, rt reflect.Type, pointerNum uint32) (any, error) { + keyType := rt.Key() + elemType := rt.Elem() + + keyTypeName, err := getTypeName(keyType) + if err != nil { + return nil, err + } + elemTypeName, err := getTypeName(elemType) + if err != nil { + return nil, err + } + + typeName := fmt.Sprintf("map[%s]%s", keyTypeName, elemTypeName) + if pointerNum > 0 { + typeName = strings.Repeat("*", int(pointerNum)) + typeName + } + + return map[string]any{ + typeFieldName: typeName, + "value": value, + }, nil +} + +func wrapSliceWithType(value []any, rt reflect.Type, pointerNum uint32) (any, error) { + elemType := rt.Elem() + elemTypeName, err := getTypeName(elemType) + if err != nil { + return nil, err + } + + // Preserve array vs slice distinction on the wire so the type round-trips + // exactly. Without this branch, a [N]T value placed in an interface field + // would silently come back as []T. + var typeName string + if rt.Kind() == reflect.Array { + typeName = fmt.Sprintf("[%d]%s", rt.Len(), elemTypeName) + } else { + typeName = fmt.Sprintf("[]%s", elemTypeName) + } + if pointerNum > 0 { + typeName = strings.Repeat("*", int(pointerNum)) + typeName + } + + return map[string]any{ + typeFieldName: typeName, + "value": value, + }, nil +} + +func getTypeName(t reflect.Type) (string, error) { + var pointerPrefix string + for t.Kind() == reflect.Ptr { + pointerPrefix += "*" + t = t.Elem() + } + + if t.Kind() == reflect.Map { + keyName, err := getTypeName(t.Key()) + if err != nil { + return "", err + } + elemName, err := getTypeName(t.Elem()) + if err != nil { + return "", err + } + return pointerPrefix + fmt.Sprintf("map[%s]%s", keyName, elemName), nil + } + + if t.Kind() == reflect.Slice { + elemName, err := getTypeName(t.Elem()) + if err != nil { + return "", err + } + return pointerPrefix + fmt.Sprintf("[]%s", elemName), nil + } + + if t.Kind() == reflect.Array { + elemName, err := getTypeName(t.Elem()) + if err != nil { + return "", err + } + return pointerPrefix + fmt.Sprintf("[%d]%s", t.Len(), elemName), nil + } + + key, ok := rm[t] + if !ok { + return "", fmt.Errorf("unknown type: %v", t) + } + return pointerPrefix + key, nil +} + +func hrUnmarshal(data any, targetType reflect.Type) (any, error) { + if data == nil { + return nil, nil + } + + ptrNum, baseType := derefPointerNum(targetType) + + switch v := data.(type) { + case map[string]any: + return hrUnmarshalMap(v, targetType, baseType, ptrNum) + case []any: + return hrUnmarshalSlice(v, targetType, baseType, ptrNum) + default: + return hrUnmarshalPrimitive(data, targetType, baseType, ptrNum) + } +} + +func hrUnmarshalMap(data map[string]any, targetType, baseType reflect.Type, ptrNum uint32) (any, error) { + if typeStr, hasType := data[typeFieldName].(string); hasType { + if shouldTreatAsTypeEnvelope(data, baseType, typeStr) { + return hrUnmarshalTyped(data, typeStr) + } + } + + if baseType.Kind() == reflect.Struct { + return hrUnmarshalStruct(data, targetType, baseType, ptrNum) + } + + if baseType.Kind() == reflect.Map { + return hrUnmarshalMapValue(data, targetType, baseType, ptrNum) + } + + if baseType.Kind() == reflect.Interface { + result := make(map[string]any) + for k, v := range data { + unmarshaled, err := hrUnmarshal(v, reflect.TypeOf((*any)(nil)).Elem()) + if err != nil { + return nil, err + } + result[k] = unmarshaled + } + return result, nil + } + + return nil, fmt.Errorf("cannot unmarshal map to %v", targetType) +} + +func shouldTreatAsTypeEnvelope(data map[string]any, baseType reflect.Type, typeStr string) bool { + if _, _, err := parseTypeName(typeStr); err != nil { + return false + } + if baseType.Kind() == reflect.Interface { + return true + } + _, hasValue := data["value"] + return hasValue && len(data) == 2 +} + +func hrUnmarshalTyped(data map[string]any, typeStr string) (any, error) { + actualType, ptrNum, err := parseTypeName(typeStr) + if err != nil { + return nil, err + } + + value, hasValue := data["value"] + if hasValue && len(data) == 2 { + result, unmarshalErr := hrUnmarshal(value, actualType) + if unmarshalErr != nil { + return nil, unmarshalErr + } + return wrapPointers(result, ptrNum), nil + } + + dataCopy := make(map[string]any) + for k, v := range data { + if k != typeFieldName { + dataCopy[k] = v + } + } + + result, err := hrUnmarshal(dataCopy, actualType) + if err != nil { + return nil, err + } + return wrapPointers(result, ptrNum), nil +} + +func hrUnmarshalStruct(data map[string]any, targetType, baseType reflect.Type, ptrNum uint32) (any, error) { + if checkMarshaler(baseType) { + jsonBytes, err := sonic.Marshal(data) + if err != nil { + return nil, fmt.Errorf("failed to marshal data for custom unmarshaler: %w", err) + } + result := reflect.New(baseType) + if err := json.Unmarshal(jsonBytes, result.Interface()); err != nil { + return nil, fmt.Errorf("failed to unmarshal with custom unmarshaler: %w", err) + } + return wrapPointers(result.Elem().Interface(), ptrNum), nil + } + + result, dResult := createValueFromType(targetType) + + for i := 0; i < baseType.NumField(); i++ { + field := baseType.Field(i) + if field.PkgPath != "" { + continue + } + + jsonTag := field.Tag.Get("json") + fieldName := getJSONFieldName(field.Name, jsonTag) + if fieldName == "-" { + continue + } + + fieldData, ok := data[fieldName] + if !ok { + fieldData, ok = data[field.Name] + } + if !ok { + continue + } + + fieldValue, err := hrUnmarshal(fieldData, field.Type) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal field %s: %w", field.Name, err) + } + + if fieldValue != nil { + fieldRef := dResult.FieldByName(field.Name) + if fieldRef.CanSet() { + if !setValueWithConversion(fieldRef, reflect.ValueOf(fieldValue)) { + return nil, fmt.Errorf("cannot set field %s: type mismatch", field.Name) + } + } + } + } + + return result.Interface(), nil +} + +func hrUnmarshalMapValue(data map[string]any, targetType, baseType reflect.Type, ptrNum uint32) (any, error) { + result, dResult := createValueFromType(targetType) + + keyType := baseType.Key() + elemType := baseType.Elem() + + for k, v := range data { + var keyValue reflect.Value + if keyType.Kind() == reflect.String { + keyValue = reflect.ValueOf(k) + } else { + keyPtr := reflect.New(keyType) + if err := sonic.UnmarshalString(k, keyPtr.Interface()); err != nil { + return nil, fmt.Errorf("failed to unmarshal map key %s: %w", k, err) + } + keyValue = keyPtr.Elem() + } + + elemValue, err := hrUnmarshal(v, elemType) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal map value for key %s: %w", k, err) + } + + if elemValue == nil { + dResult.SetMapIndex(keyValue, reflect.Zero(elemType)) + } else { + dResult.SetMapIndex(keyValue, reflect.ValueOf(elemValue)) + } + } + + return result.Interface(), nil +} + +func hrUnmarshalSlice(data []any, targetType, baseType reflect.Type, ptrNum uint32) (any, error) { + if baseType.Kind() == reflect.Interface { + result := make([]any, len(data)) + for i, elem := range data { + unmarshaled, err := hrUnmarshal(elem, reflect.TypeOf((*any)(nil)).Elem()) + if err != nil { + return nil, err + } + result[i] = unmarshaled + } + return result, nil + } + + if baseType.Kind() != reflect.Slice && baseType.Kind() != reflect.Array { + return nil, fmt.Errorf("cannot unmarshal slice to %v", targetType) + } + + elemType := baseType.Elem() + result, dResult := createValueFromType(targetType) + + if baseType.Kind() == reflect.Array { + for i, elem := range data { + if i >= dResult.Len() { + break + } + elemValue, err := hrUnmarshal(elem, elemType) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal array element %d: %w", i, err) + } + if elemValue == nil { + dResult.Index(i).Set(reflect.Zero(elemType)) + } else { + dResult.Index(i).Set(reflect.ValueOf(elemValue)) + } + } + } else { + for i, elem := range data { + elemValue, err := hrUnmarshal(elem, elemType) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal slice element %d: %w", i, err) + } + if elemValue == nil { + dResult.Set(reflect.Append(dResult, reflect.Zero(elemType))) + } else { + dResult.Set(reflect.Append(dResult, reflect.ValueOf(elemValue))) + } + } + } + + return result.Interface(), nil +} + +func hrUnmarshalPrimitive(data any, targetType, baseType reflect.Type, ptrNum uint32) (any, error) { + if baseType.Kind() == reflect.Interface { + return convertJSONPrimitive(data), nil + } + + if n, ok := data.(json.Number); ok { + switch baseType.Kind() { + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + i, err := strconv.ParseInt(n.String(), 10, baseType.Bits()) + if err != nil { + return nil, fmt.Errorf("failed to parse %q as %v: %w", n.String(), baseType, err) + } + result := reflect.New(baseType).Elem() + result.SetInt(i) + return wrapPointers(result.Interface(), ptrNum), nil + case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr: + u, err := strconv.ParseUint(n.String(), 10, baseType.Bits()) + if err != nil { + return nil, fmt.Errorf("failed to parse %q as %v: %w", n.String(), baseType, err) + } + result := reflect.New(baseType).Elem() + result.SetUint(u) + return wrapPointers(result.Interface(), ptrNum), nil + case reflect.Float32, reflect.Float64: + f, err := strconv.ParseFloat(n.String(), baseType.Bits()) + if err != nil { + return nil, fmt.Errorf("failed to parse %q as %v: %w", n.String(), baseType, err) + } + result := reflect.New(baseType).Elem() + result.SetFloat(f) + return wrapPointers(result.Interface(), ptrNum), nil + } + } + + dataValue := reflect.ValueOf(data) + if dataValue.Type().AssignableTo(baseType) { + return wrapPointers(data, ptrNum), nil + } + + if dataValue.Type().ConvertibleTo(baseType) { + converted := dataValue.Convert(baseType) + return wrapPointers(converted.Interface(), ptrNum), nil + } + + if baseType.Kind() == reflect.Int || baseType.Kind() == reflect.Int8 || + baseType.Kind() == reflect.Int16 || baseType.Kind() == reflect.Int32 || + baseType.Kind() == reflect.Int64 { + if f, ok := data.(float64); ok { + result := reflect.New(baseType).Elem() + result.SetInt(int64(f)) + return wrapPointers(result.Interface(), ptrNum), nil + } + } + + if baseType.Kind() == reflect.Uint || baseType.Kind() == reflect.Uint8 || + baseType.Kind() == reflect.Uint16 || baseType.Kind() == reflect.Uint32 || + baseType.Kind() == reflect.Uint64 { + if f, ok := data.(float64); ok { + result := reflect.New(baseType).Elem() + result.SetUint(uint64(f)) + return wrapPointers(result.Interface(), ptrNum), nil + } + } + + if baseType.Kind() == reflect.Float32 { + if f, ok := data.(float64); ok { + return wrapPointers(float32(f), ptrNum), nil + } + } + + jsonBytes, err := sonic.Marshal(data) + if err != nil { + return nil, fmt.Errorf("failed to re-marshal data: %w", err) + } + + result := reflect.New(baseType) + if err := sonic.Unmarshal(jsonBytes, result.Interface()); err != nil { + return nil, fmt.Errorf("failed to unmarshal to %v: %w", baseType, err) + } + + return wrapPointers(result.Elem().Interface(), ptrNum), nil +} + +func parseTypeName(typeStr string) (reflect.Type, uint32, error) { + var ptrNum uint32 + for strings.HasPrefix(typeStr, "*") { + ptrNum++ + typeStr = typeStr[1:] + } + + if strings.HasPrefix(typeStr, "map[") { + return parseMapType(typeStr, ptrNum) + } + + if strings.HasPrefix(typeStr, "[]") { + return parseSliceType(typeStr, ptrNum) + } + + if strings.HasPrefix(typeStr, "[") { + return parseArrayType(typeStr, ptrNum) + } + + rt, ok := m[typeStr] + if !ok { + return nil, 0, fmt.Errorf("unknown type: %s", typeStr) + } + + return rt, ptrNum, nil +} + +func parseMapType(typeStr string, ptrNum uint32) (reflect.Type, uint32, error) { + inner := typeStr[4:] + bracketCount := 1 + keyEnd := 0 + for i, c := range inner { + if c == '[' { + bracketCount++ + } else if c == ']' { + bracketCount-- + if bracketCount == 0 { + keyEnd = i + break + } + } + } + + keyTypeStr := inner[:keyEnd] + valueTypeStr := inner[keyEnd+1:] + + keyType, keyPtrNum, err := parseTypeName(keyTypeStr) + if err != nil { + return nil, 0, fmt.Errorf("failed to parse map key type: %w", err) + } + + finalKeyType := keyType + for i := uint32(0); i < keyPtrNum; i++ { + finalKeyType = reflect.PointerTo(finalKeyType) + } + + valueType, valuePtrNum, err := parseTypeName(valueTypeStr) + if err != nil { + return nil, 0, fmt.Errorf("failed to parse map value type: %w", err) + } + + finalValueType := valueType + for i := uint32(0); i < valuePtrNum; i++ { + finalValueType = reflect.PointerTo(finalValueType) + } + + return reflect.MapOf(finalKeyType, finalValueType), ptrNum, nil +} + +func parseSliceType(typeStr string, ptrNum uint32) (reflect.Type, uint32, error) { + elemTypeStr := typeStr[2:] + elemType, elemPtrNum, err := parseTypeName(elemTypeStr) + if err != nil { + return nil, 0, fmt.Errorf("failed to parse slice element type: %w", err) + } + + finalElemType := elemType + for i := uint32(0); i < elemPtrNum; i++ { + finalElemType = reflect.PointerTo(finalElemType) + } + + return reflect.SliceOf(finalElemType), ptrNum, nil +} + +func parseArrayType(typeStr string, ptrNum uint32) (reflect.Type, uint32, error) { + closeBracket := strings.Index(typeStr, "]") + if closeBracket == -1 { + return nil, 0, fmt.Errorf("invalid array type: %s", typeStr) + } + + sizeStr := typeStr[1:closeBracket] + var size int + if _, err := fmt.Sscanf(sizeStr, "%d", &size); err != nil { + return nil, 0, fmt.Errorf("invalid array size: %s", sizeStr) + } + + elemTypeStr := typeStr[closeBracket+1:] + elemType, elemPtrNum, err := parseTypeName(elemTypeStr) + if err != nil { + return nil, 0, fmt.Errorf("failed to parse array element type: %w", err) + } + + finalElemType := elemType + for i := uint32(0); i < elemPtrNum; i++ { + finalElemType = reflect.PointerTo(finalElemType) + } + + return reflect.ArrayOf(size, finalElemType), ptrNum, nil +} + +func wrapPointers(value any, ptrNum uint32) any { + if ptrNum == 0 || value == nil { + return value + } + + rv := reflect.ValueOf(value) + for i := uint32(0); i < ptrNum; i++ { + ptr := reflect.New(rv.Type()) + ptr.Elem().Set(rv) + rv = ptr + } + return rv.Interface() +} + +func convertJSONPrimitive(data any) any { + switch v := data.(type) { + case json.Number: + s := v.String() + if strings.ContainsAny(s, ".eE") { + if f, err := strconv.ParseFloat(s, 64); err == nil { + return f + } + return s + } + if i, err := strconv.ParseInt(s, 10, 0); err == nil { + return int(i) + } + if u, err := strconv.ParseUint(s, 10, 64); err == nil { + return u + } + return s + case float64: + if v == float64(int64(v)) { + return int(v) + } + return v + default: + return data + } +} + +func setValueWithConversion(target, source reflect.Value) bool { + if !source.IsValid() { + target.Set(reflect.Zero(target.Type())) + return true + } + + if source.Type().AssignableTo(target.Type()) { + target.Set(source) + return true + } + + if target.Kind() == reflect.Ptr { + if target.IsNil() && target.CanSet() { + target.Set(reflect.New(target.Type().Elem())) + } + return setValueWithConversion(target.Elem(), source) + } + + if source.Kind() == reflect.Ptr { + if source.IsNil() { + target.Set(reflect.Zero(target.Type())) + return true + } + return setValueWithConversion(target, source.Elem()) + } + + if source.Type().ConvertibleTo(target.Type()) { + target.Set(source.Convert(target.Type())) + return true + } + + if target.Kind() == reflect.Int || target.Kind() == reflect.Int8 || + target.Kind() == reflect.Int16 || target.Kind() == reflect.Int32 || + target.Kind() == reflect.Int64 { + if source.Kind() == reflect.Float64 { + target.SetInt(int64(source.Float())) + return true + } + if source.Kind() == reflect.Int { + target.SetInt(int64(source.Int())) + return true + } + } + + if target.Kind() == reflect.Uint || target.Kind() == reflect.Uint8 || + target.Kind() == reflect.Uint16 || target.Kind() == reflect.Uint32 || + target.Kind() == reflect.Uint64 { + if source.Kind() == reflect.Float64 { + target.SetUint(uint64(source.Float())) + return true + } + } + + if target.Kind() == reflect.Float32 || target.Kind() == reflect.Float64 { + if source.Kind() == reflect.Float64 { + target.SetFloat(source.Float()) + return true + } + if source.Kind() == reflect.Int { + target.SetFloat(float64(source.Int())) + return true + } + } + + return false +} + +func decodeJSONAny(data []byte) (any, error) { + dec := json.NewDecoder(bytes.NewReader(data)) + dec.UseNumber() + var raw any + if err := dec.Decode(&raw); err != nil { + return nil, err + } + return raw, nil +} + +func getJSONFieldName(fieldName, jsonTag string) string { + if jsonTag == "" { + return fieldName + } + parts := strings.Split(jsonTag, ",") + if parts[0] == "" { + return fieldName + } + return parts[0] +} + +func hasOmitempty(jsonTag string) bool { + return strings.Contains(jsonTag, "omitempty") +} + +func isEmptyValue(v reflect.Value) bool { + switch v.Kind() { + case reflect.Array, reflect.Map, reflect.Slice, reflect.String: + return v.Len() == 0 + case reflect.Bool: + return !v.Bool() + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + return v.Int() == 0 + case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr: + return v.Uint() == 0 + case reflect.Float32, reflect.Float64: + return v.Float() == 0 + case reflect.Interface, reflect.Ptr: + return v.IsNil() + } + return false +} + +func isPrimitiveJSONType(t reflect.Type) bool { + switch t.Kind() { + case reflect.Bool, reflect.String, + reflect.Float32, reflect.Float64: + return true + default: + return false + } +} diff --git a/internal/serialization/human_readable_test.go b/internal/serialization/human_readable_test.go new file mode 100644 index 000000000..caf2205fb --- /dev/null +++ b/internal/serialization/human_readable_test.go @@ -0,0 +1,1482 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package serialization + +import ( + "encoding/json" + "reflect" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ===== Mock type replacing schema.ToolInfo (pointer-receiver MarshalJSON) ===== + +type hrMockToolInfo struct { + Name string + Desc string + params map[string]string // unexported — only via MarshalJSON +} + +type hrMockToolInfoJSON struct { + Name string `json:"name"` + Desc string `json:"desc"` + HasParams bool `json:"has_params"` + Params map[string]string `json:"params,omitempty"` +} + +func (t *hrMockToolInfo) MarshalJSON() ([]byte, error) { + tmp := &hrMockToolInfoJSON{Name: t.Name, Desc: t.Desc} + if t.params != nil { + tmp.HasParams = true + tmp.Params = t.params + } + return json.Marshal(tmp) +} + +func (t *hrMockToolInfo) UnmarshalJSON(data []byte) error { + tmp := &hrMockToolInfoJSON{} + if err := json.Unmarshal(data, tmp); err != nil { + return err + } + t.Name = tmp.Name + t.Desc = tmp.Desc + if tmp.HasParams { + t.params = tmp.Params + } + return nil +} + +func newMockToolInfo(name, desc string, params map[string]string) *hrMockToolInfo { + return &hrMockToolInfo{Name: name, Desc: desc, params: params} +} + +// Holder types for position tests. +type hrMockToolInfoConcreteHolder struct { + T *hrMockToolInfo `json:"t"` +} +type hrMockToolInfoInterfaceHolder struct { + V any `json:"v"` +} +type hrMockToolInfoSliceHolder struct { + S []*hrMockToolInfo `json:"s"` +} +type hrMockToolInfoMapHolder struct { + M map[string]*hrMockToolInfo `json:"m"` +} + +// ===== Existing fixture types ===== + +type hrTestStruct struct { + Name string `json:"name"` + Value int `json:"value"` +} + +type hrTestStructWithExtra struct { + Name string `json:"name"` + Extra map[string]any `json:"extra,omitempty"` +} + +type hrStructWithInterface struct { + A any + B any + C map[string]any +} + +type hrWrapper struct { + Inner hrTestStruct `json:"inner"` +} + +type hrLargeIntegerStruct struct { + I int64 `json:"i"` + U uint64 `json:"u"` + A any `json:"a"` +} + +type hrReservedTypeStruct struct { + Type string `json:"$type"` + Name string `json:"name"` +} + +type hrZeroValueStruct struct { + S string `json:"s"` + I int `json:"i"` + B bool `json:"b"` +} + +// ===== Edge-case fixture types ===== + +type hrEdgeArrayHolder struct { + A [3]int `json:"a"` + B [2]string `json:"b"` + I any `json:"i"` +} + +type hrEdgeIntKeyMap struct { + M map[int]string `json:"m"` +} + +type hrEdgeStructKeyMap struct { + M map[hrEdgeKey]string `json:"m"` +} + +type hrEdgeKey struct { + K1 string `json:"k1"` + K2 int `json:"k2"` +} + +type hrEdgePtrLevels struct { + P *int `json:"p"` + Q **int `json:"q"` + R ***int `json:"r"` +} + +type hrEdgeNestedSlicePtr struct { + S []*hrEdgeAtom `json:"s"` + M map[string]*hrEdgeAtom `json:"m"` +} + +type hrEdgeAtom struct { + N int `json:"n"` +} + +type hrEdgeNumericConvert struct { + I8 int8 `json:"i8"` + I16 int16 `json:"i16"` + I32 int32 `json:"i32"` + U8 uint8 `json:"u8"` + U16 uint16 `json:"u16"` + U32 uint32 `json:"u32"` + F32 float32 `json:"f32"` +} + +type hrEdgeFancyJSON struct { + V hrJSONMarshaler `json:"v"` +} + +type hrJSONMarshaler struct { + Inner string +} + +func (m hrJSONMarshaler) MarshalJSON() ([]byte, error) { + return []byte(`"prefix:` + m.Inner + `"`), nil +} + +func (m *hrJSONMarshaler) UnmarshalJSON(data []byte) error { + s := strings.Trim(string(data), `"`) + m.Inner = strings.TrimPrefix(s, "prefix:") + return nil +} + +type hrEdgeIgnoreField struct { + A string `json:"a"` + B string `json:"-"` + C string + d string //nolint:unused // intentional: unexported field probes filtering +} + +type hrEdgeAnyContainer struct { + V any `json:"v"` +} + +type hrEdgeUnregisteredField struct { + V hrUnregisteredInner `json:"v"` +} + +// hrUnregisteredInner is intentionally NOT passed to GenericRegister so we can +// observe how the serializer treats concrete-typed (non-interface) fields whose +// type isn't registered. +type hrUnregisteredInner struct { + N int `json:"n"` +} + +// hrUnregisteredHere is intentionally never registered. Used only in +// TestHumanReadableSerializer_MarshalErrors. +type hrUnregisteredHere struct { + X int +} + +// ===== init: type registrations ===== + +func init() { + // Basic fixture types. + _ = GenericRegister[hrTestStruct]("hr_test_struct") + _ = GenericRegister[hrTestStructWithExtra]("hr_test_struct_with_extra") + _ = GenericRegister[hrStructWithInterface]("hr_struct_with_interface") + _ = GenericRegister[hrWrapper]("hr_wrapper") + _ = GenericRegister[hrLargeIntegerStruct]("hr_large_integer_struct") + _ = GenericRegister[hrReservedTypeStruct]("hr_reserved_type_struct") + _ = GenericRegister[hrZeroValueStruct]("hr_zero_value_struct") + + // Edge-case types. + _ = GenericRegister[hrEdgeArrayHolder]("hr_edge_array_holder") + _ = GenericRegister[[3]int]("hr_edge_array_3_int") + _ = GenericRegister[[2]string]("hr_edge_array_2_string") + _ = GenericRegister[hrEdgeIntKeyMap]("hr_edge_int_key_map") + _ = GenericRegister[hrEdgeStructKeyMap]("hr_edge_struct_key_map") + _ = GenericRegister[hrEdgeKey]("hr_edge_key") + _ = GenericRegister[hrEdgePtrLevels]("hr_edge_ptr_levels") + _ = GenericRegister[hrEdgeNestedSlicePtr]("hr_edge_nested_slice_ptr") + _ = GenericRegister[hrEdgeAtom]("hr_edge_atom") + _ = GenericRegister[hrEdgeNumericConvert]("hr_edge_numeric_convert") + _ = GenericRegister[hrEdgeFancyJSON]("hr_edge_fancy_json") + _ = GenericRegister[hrJSONMarshaler]("hr_edge_json_marshaler") + _ = GenericRegister[hrEdgeIgnoreField]("hr_edge_ignore_field") + _ = GenericRegister[hrEdgeAnyContainer]("hr_edge_any_container") + + // Mock ToolInfo types. + _ = GenericRegister[hrMockToolInfo]("hr_mock_tool_info") + _ = GenericRegister[hrMockToolInfoConcreteHolder]("hr_mock_tool_info_concrete_holder") + _ = GenericRegister[hrMockToolInfoInterfaceHolder]("hr_mock_tool_info_interface_holder") + _ = GenericRegister[hrMockToolInfoSliceHolder]("hr_mock_tool_info_slice_holder") + _ = GenericRegister[hrMockToolInfoMapHolder]("hr_mock_tool_info_map_holder") +} + +// ============================================================================= +// Section: Basic serialization behavior +// ============================================================================= + +func TestHumanReadableSerializer_OmitemptyBehavior(t *testing.T) { + s := &HumanReadableSerializer{} + + input := hrTestStructWithExtra{ + Name: "test", + Extra: nil, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var jsonMap map[string]any + err = json.Unmarshal(data, &jsonMap) + require.NoError(t, err) + + _, hasExtra := jsonMap["extra"] + assert.False(t, hasExtra, "omitempty field should not be present when nil") +} + +func TestHumanReadableSerializer_JSONFieldNames(t *testing.T) { + s := &HumanReadableSerializer{} + + input := hrTestStruct{ + Name: "test", + Value: 123, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var jsonMap map[string]any + err = json.Unmarshal(data, &jsonMap) + require.NoError(t, err) + + assert.Equal(t, "test", jsonMap["name"]) + assert.Equal(t, float64(123), jsonMap["value"]) + _, hasName := jsonMap["Name"] + assert.False(t, hasName, "should use json tag name, not struct field name") +} + +func TestHumanReadableSerializer_NonOmitEmptyZeroValuesAreScalars(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrZeroValueStruct{} + + data, err := s.Marshal(input) + require.NoError(t, err) + + var raw map[string]any + err = json.Unmarshal(data, &raw) + require.NoError(t, err) + assert.Equal(t, "", raw["s"]) + assert.Equal(t, float64(0), raw["i"]) + assert.Equal(t, false, raw["b"]) + + var result hrZeroValueStruct + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Equal(t, input, result) +} + +func TestHumanReadableSerializer_CompareWithInternalSerializer(t *testing.T) { + hr := &HumanReadableSerializer{} + is := &InternalSerializer{} + + input := hrStructWithInterface{ + A: "string", + B: hrTestStruct{Name: "test", Value: 42}, + C: map[string]any{ + "key1": "value1", + "key2": 123, + }, + } + + hrData, err := hr.Marshal(input) + require.NoError(t, err) + + isData, err := is.Marshal(input) + require.NoError(t, err) + + t.Logf("HumanReadable output size: %d bytes", len(hrData)) + t.Logf("Internal output size: %d bytes", len(isData)) + t.Logf("HumanReadable output:\n%s", string(hrData)) + + assert.Less(t, len(hrData), len(isData), "HumanReadable should produce smaller output") + + var hrResult hrStructWithInterface + err = hr.Unmarshal(hrData, &hrResult) + require.NoError(t, err) + + var isResult hrStructWithInterface + err = is.Unmarshal(isData, &isResult) + require.NoError(t, err) + + assert.Equal(t, hrResult.A, isResult.A) + assert.Equal(t, hrResult.B, isResult.B) +} + +// ============================================================================= +// Section: Type annotations ($type envelope) +// ============================================================================= + +func TestHumanReadableSerializer_TypeAnnotationOnlyForInterfaceFields(t *testing.T) { + s := &HumanReadableSerializer{} + + t.Run("concrete struct field has no $type", func(t *testing.T) { + input := hrWrapper{ + Inner: hrTestStruct{Name: "test", Value: 123}, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var jsonMap map[string]any + err = json.Unmarshal(data, &jsonMap) + require.NoError(t, err) + + innerMap := jsonMap["inner"].(map[string]any) + _, hasType := innerMap["$type"] + assert.False(t, hasType, "concrete struct field should not have $type annotation") + }) + + t.Run("interface field has $type", func(t *testing.T) { + input := hrStructWithInterface{ + A: hrTestStruct{Name: "test", Value: 123}, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var jsonMap map[string]any + err = json.Unmarshal(data, &jsonMap) + require.NoError(t, err) + + aMap := jsonMap["A"].(map[string]any) + _, hasType := aMap["$type"] + assert.True(t, hasType, "interface field should have $type annotation") + }) +} + +func TestHumanReadableSerializer_PreservesUserTypeKey(t *testing.T) { + s := &HumanReadableSerializer{} + + t.Run("map key", func(t *testing.T) { + input := map[string]any{ + "$type": "user-controlled", + "value": int64(7), + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var result map[string]any + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Equal(t, input, result) + }) + + t.Run("struct field", func(t *testing.T) { + input := hrReservedTypeStruct{ + Type: "user-controlled", + Name: "kept", + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var result hrReservedTypeStruct + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Equal(t, input, result) + }) + + t.Run("struct field with registered value", func(t *testing.T) { + input := hrReservedTypeStruct{ + Type: "_eino_string", + Name: "kept", + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var result hrReservedTypeStruct + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Equal(t, input, result) + }) +} + +// ============================================================================= +// Section: Pointer-receiver MarshalJSON (mock ToolInfo regression) +// ============================================================================= + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_TopLevel(t *testing.T) { + s := &HumanReadableSerializer{} + + original := newMockToolInfo("search", "search the docs", map[string]string{"q": "query"}) + + data, err := s.Marshal(original) + require.NoError(t, err) + + // The wire format must come from MarshalJSON (lowercase tags from hrMockToolInfoJSON). + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + assert.Equal(t, "search", raw["name"], "must use MarshalJSON's lowercase 'name' tag") + assert.Equal(t, "search the docs", raw["desc"]) + assert.Equal(t, true, raw["has_params"], "MarshalJSON must record has_params=true") + require.Contains(t, raw, "params", "MarshalJSON must include params") + + var got hrMockToolInfo + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, original.Name, got.Name) + assert.Equal(t, original.Desc, got.Desc) + assert.Equal(t, original.params, got.params, "unexported params must round-trip via MarshalJSON/UnmarshalJSON") +} + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_ValueType(t *testing.T) { + s := &HumanReadableSerializer{} + + // Pass by value (not pointer) — exercises the addressability shim in hrMarshalStruct. + tiVal := hrMockToolInfo{ + Name: "value-type", + Desc: "no pointer", + params: map[string]string{"x": "y"}, + } + + data, err := s.Marshal(tiVal) + require.NoError(t, err) + + // The wire format must come from MarshalJSON (lowercase tags). + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + assert.Equal(t, "value-type", raw["name"], + "value-type hrMockToolInfo must still go through pointer-receiver MarshalJSON via the addressability shim") + assert.Equal(t, true, raw["has_params"]) + + // Round-trip into a value target. + var got hrMockToolInfo + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, "value-type", got.Name) + assert.Equal(t, map[string]string{"x": "y"}, got.params) +} + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_NoParams(t *testing.T) { + s := &HumanReadableSerializer{} + + original := newMockToolInfo("ping", "no-arg tool", nil) + + data, err := s.Marshal(original) + require.NoError(t, err) + + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + assert.Equal(t, false, raw["has_params"], "nil params → has_params=false") + _, hasParams := raw["params"] + assert.False(t, hasParams, "nil params → params field must be absent (omitempty)") + + var got hrMockToolInfo + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, "ping", got.Name) + assert.Equal(t, "no-arg tool", got.Desc) + assert.Nil(t, got.params, "absent params must remain nil") +} + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_InConcreteField(t *testing.T) { + s := &HumanReadableSerializer{} + holder := hrMockToolInfoConcreteHolder{ + T: newMockToolInfo("search", "search docs", map[string]string{"q": "query"}), + } + + data, err := s.Marshal(holder) + require.NoError(t, err) + + // Concrete fields don't carry a $type envelope. + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + tMap, ok := raw["t"].(map[string]any) + require.True(t, ok) + _, hasType := tMap["$type"] + assert.False(t, hasType, "concrete pointer field should not carry a $type envelope") + assert.Equal(t, "search", tMap["name"], "must still go through MarshalJSON") + + var got hrMockToolInfoConcreteHolder + require.NoError(t, s.Unmarshal(data, &got)) + require.NotNil(t, got.T) + assert.Equal(t, holder.T.Name, got.T.Name) + assert.Equal(t, holder.T.params, got.T.params) +} + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_InInterfaceField(t *testing.T) { + s := &HumanReadableSerializer{} + holder := hrMockToolInfoInterfaceHolder{ + V: newMockToolInfo("search", "in interface", map[string]string{"q": "query"}), + } + + data, err := s.Marshal(holder) + require.NoError(t, err) + + // Interface fields must include the $type envelope. + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + vMap, ok := raw["v"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "*hr_mock_tool_info", vMap["$type"], + "interface field with *hrMockToolInfo must carry the registered type tag") + + var got hrMockToolInfoInterfaceHolder + require.NoError(t, s.Unmarshal(data, &got)) + + gotTI, ok := got.V.(*hrMockToolInfo) + require.True(t, ok, "interface field must reconstruct as *hrMockToolInfo, got %T", got.V) + assert.Equal(t, "search", gotTI.Name) + assert.Equal(t, map[string]string{"q": "query"}, gotTI.params, "params must round-trip through interface field") +} + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_InSlice(t *testing.T) { + s := &HumanReadableSerializer{} + holder := hrMockToolInfoSliceHolder{ + S: []*hrMockToolInfo{ + newMockToolInfo("t1", "first", nil), + newMockToolInfo("t2", "second", map[string]string{"x": "1"}), + nil, // nil pointer in slice — must round-trip as nil. + }, + } + + data, err := s.Marshal(holder) + require.NoError(t, err) + + var got hrMockToolInfoSliceHolder + require.NoError(t, s.Unmarshal(data, &got)) + require.Len(t, got.S, 3) + require.NotNil(t, got.S[0]) + assert.Equal(t, "t1", got.S[0].Name) + assert.Nil(t, got.S[0].params) + require.NotNil(t, got.S[1]) + assert.Equal(t, map[string]string{"x": "1"}, got.S[1].params) + assert.Nil(t, got.S[2], "nil entry in slice must round-trip as nil") +} + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_InMap(t *testing.T) { + s := &HumanReadableSerializer{} + holder := hrMockToolInfoMapHolder{ + M: map[string]*hrMockToolInfo{ + "alpha": newMockToolInfo("alpha", "first", nil), + "beta": newMockToolInfo("beta", "second", map[string]string{"y": "2"}), + "nilEntry": nil, + }, + } + + data, err := s.Marshal(holder) + require.NoError(t, err) + + var got hrMockToolInfoMapHolder + require.NoError(t, s.Unmarshal(data, &got)) + require.Len(t, got.M, 3) + require.NotNil(t, got.M["alpha"]) + assert.Equal(t, "alpha", got.M["alpha"].Name) + require.NotNil(t, got.M["beta"]) + assert.Equal(t, map[string]string{"y": "2"}, got.M["beta"].params) + assert.Nil(t, got.M["nilEntry"]) +} + +func TestHumanReadableSerializer_PtrReceiverMarshalJSON_NilPointer(t *testing.T) { + s := &HumanReadableSerializer{} + var nilTI *hrMockToolInfo + + data, err := s.Marshal(nilTI) + require.NoError(t, err) + assert.Equal(t, "null", string(data), "nil pointer must marshal to JSON null") + + var got *hrMockToolInfo + require.NoError(t, s.Unmarshal(data, &got)) + assert.Nil(t, got) +} + +// ============================================================================= +// Section: Fixed-size arrays +// ============================================================================= + +func TestHumanReadableSerializer_FixedSizeArrayRoundTrip(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeArrayHolder{ + A: [3]int{10, 20, 30}, + B: [2]string{"x", "y"}, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got hrEdgeArrayHolder + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input, got) +} + +func TestHumanReadableSerializer_ArrayInInterfaceField(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeArrayHolder{ + I: [3]int{1, 2, 3}, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + // Verify $type annotation includes the array shape. + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + iMap, ok := raw["i"].(map[string]any) + require.True(t, ok, "interface field must serialize with type envelope") + require.Contains(t, iMap, "$type") + assert.Contains(t, iMap["$type"], "[3]") + + var got hrEdgeArrayHolder + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input.I, got.I) +} + +func TestHumanReadableSerializer_ArrayWithExtraJSONElementsTruncates(t *testing.T) { + s := &HumanReadableSerializer{} + + original := hrEdgeArrayHolder{A: [3]int{1, 2, 3}} + data, err := s.Marshal(original) + require.NoError(t, err) + + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + raw["a"] = []any{json.Number("11"), json.Number("22"), json.Number("33"), json.Number("44"), json.Number("55")} + + tampered, err := json.Marshal(raw) + require.NoError(t, err) + + var got hrEdgeArrayHolder + require.NoError(t, s.Unmarshal(tampered, &got)) + assert.Equal(t, [3]int{11, 22, 33}, got.A, + "extra JSON elements beyond array length must be silently dropped") +} + +// ============================================================================= +// Section: Non-string map keys +// ============================================================================= + +func TestHumanReadableSerializer_IntegerMapKeys(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeIntKeyMap{M: map[int]string{1: "one", 2: "two", 42: "forty-two"}} + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got hrEdgeIntKeyMap + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input, got) +} + +func TestHumanReadableSerializer_StructMapKeys(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeStructKeyMap{ + M: map[hrEdgeKey]string{ + {K1: "alpha", K2: 1}: "first", + {K1: "beta", K2: 2}: "second", + }, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got hrEdgeStructKeyMap + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input, got) +} + +// ============================================================================= +// Section: Pointer indirection +// ============================================================================= + +func TestHumanReadableSerializer_MultiLevelPointers(t *testing.T) { + s := &HumanReadableSerializer{} + v1 := 7 + pv1 := &v1 + ppv1 := &pv1 + input := hrEdgePtrLevels{P: &v1, Q: &pv1, R: &ppv1} + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got hrEdgePtrLevels + require.NoError(t, s.Unmarshal(data, &got)) + require.NotNil(t, got.P) + require.NotNil(t, got.Q) + require.NotNil(t, got.R) + assert.Equal(t, 7, *got.P) + assert.Equal(t, 7, **got.Q) + assert.Equal(t, 7, ***got.R) +} + +func TestHumanReadableSerializer_NilPointerFieldIsAbsent(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgePtrLevels{P: nil, Q: nil, R: nil} + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got hrEdgePtrLevels + require.NoError(t, s.Unmarshal(data, &got)) + assert.Nil(t, got.P) + assert.Nil(t, got.Q) + assert.Nil(t, got.R) +} + +func TestHumanReadableSerializer_SliceAndMapOfPointers(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeNestedSlicePtr{ + S: []*hrEdgeAtom{{N: 1}, nil, {N: 3}}, + M: map[string]*hrEdgeAtom{ + "a": {N: 10}, + "b": nil, + }, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got hrEdgeNestedSlicePtr + require.NoError(t, s.Unmarshal(data, &got)) + require.Equal(t, len(input.S), len(got.S)) + for i := range input.S { + if input.S[i] == nil { + assert.Nil(t, got.S[i], "nil slice element[%d] must round-trip as nil", i) + } else { + require.NotNil(t, got.S[i]) + assert.Equal(t, *input.S[i], *got.S[i]) + } + } + require.Equal(t, len(input.M), len(got.M)) + require.NotNil(t, got.M["a"]) + assert.Equal(t, 10, got.M["a"].N) + assert.Nil(t, got.M["b"], "nil map value must round-trip as nil") +} + +// ============================================================================= +// Section: Numeric types +// ============================================================================= + +func TestHumanReadableSerializer_IntegerExtremes(t *testing.T) { + type holder struct { + MinI64 int64 `json:"min_i64"` + MaxI64 int64 `json:"max_i64"` + MaxU64 uint64 `json:"max_u64"` + AnyI64 any `json:"any_i64"` + AnyU64 any `json:"any_u64"` + } + _ = GenericRegister[holder]("hr_edge_int_extremes_holder") + + s := &HumanReadableSerializer{} + input := holder{ + MinI64: -1 << 63, + MaxI64: 1<<63 - 1, + MaxU64: ^uint64(0), + AnyI64: int64(-1 << 62), + AnyU64: uint64(1<<63 + 1), + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got holder + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input.MinI64, got.MinI64) + assert.Equal(t, input.MaxI64, got.MaxI64) + assert.Equal(t, input.MaxU64, got.MaxU64) + assert.Equal(t, input.AnyI64, got.AnyI64) + assert.Equal(t, input.AnyU64, got.AnyU64) +} + +func TestHumanReadableSerializer_NumericFieldTypesRoundTrip(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeNumericConvert{ + I8: -8, + I16: -16, + I32: -32, + U8: 8, + U16: 16, + U32: 32, + F32: 1.5, + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got hrEdgeNumericConvert + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input, got) +} + +func TestHumanReadableSerializer_NumericOverflowDecodeError(t *testing.T) { + s := &HumanReadableSerializer{} + + // 200 doesn't fit in int8 (-128..127). + tampered := []byte(`{"i8":200,"i16":0,"i32":0,"u8":0,"u16":0,"u32":0,"f32":0}`) + + var got hrEdgeNumericConvert + err := s.Unmarshal(tampered, &got) + require.Error(t, err, "must reject numeric overflow rather than silently truncating") + assert.Contains(t, err.Error(), "I8") + assert.Contains(t, err.Error(), "200") +} + +func TestHumanReadableSerializer_HighPrecisionFloats(t *testing.T) { + type holder struct { + F float64 `json:"f"` + F2 float64 `json:"f2"` + A any `json:"a"` + } + _ = GenericRegister[holder]("hr_edge_precision_floats_holder") + + s := &HumanReadableSerializer{} + input := holder{ + F: 1.7976931348623157e+308, // near math.MaxFloat64 + F2: 5.0e-324, // near smallest positive subnormal + A: float64(3.141592653589793), + } + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got holder + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input.F, got.F) + assert.Equal(t, input.F2, got.F2) + assert.Equal(t, input.A, got.A) +} + +// ============================================================================= +// Section: Interface field primitives +// ============================================================================= + +func TestHumanReadableSerializer_AnyFieldPrimitives(t *testing.T) { + s := &HumanReadableSerializer{} + + cases := []struct { + name string + raw string + expected any + }{ + {"int via json.Number", `{"v":42}`, int(42)}, + {"float via json.Number", `{"v":3.14}`, 3.14}, + {"exponent float", `{"v":1e2}`, 100.0}, + {"large uint via json.Number", `{"v":18446744073709551610}`, uint64(18446744073709551610)}, + {"string", `{"v":"hello"}`, "hello"}, + {"bool true", `{"v":true}`, true}, + {"bool false", `{"v":false}`, false}, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + var got hrEdgeAnyContainer + require.NoError(t, s.Unmarshal([]byte(c.raw), &got)) + assert.Equal(t, c.expected, got.V, "raw=%s", c.raw) + }) + } +} + +func TestHumanReadableSerializer_ConvertJSONPrimitive_DefaultBranch(t *testing.T) { + // A bool reaches convertJSONPrimitive's default arm. + assert.Equal(t, true, convertJSONPrimitive(true)) + assert.Equal(t, "abc", convertJSONPrimitive("abc")) + // A nil reaches the default arm too. + assert.Equal(t, nil, convertJSONPrimitive(nil)) + // A non-integer float64 must round-trip as float64. + assert.Equal(t, 3.5, convertJSONPrimitive(float64(3.5))) + // An integer-valued float64 collapses to int. + assert.Equal(t, int(7), convertJSONPrimitive(float64(7))) +} + +// ============================================================================= +// Section: Custom MarshalJSON (simple value-type marshaler) +// ============================================================================= + +func TestHumanReadableSerializer_CustomJSONMarshaler(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeFancyJSON{V: hrJSONMarshaler{Inner: "hello"}} + + data, err := s.Marshal(input) + require.NoError(t, err) + + // The inner value should serialize as the marshaler's chosen output. + assert.Contains(t, string(data), "prefix:hello") + + var got hrEdgeFancyJSON + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input, got) +} + +// ============================================================================= +// Section: Field handling +// ============================================================================= + +func TestHumanReadableSerializer_JSONDashAndUnexportedFields(t *testing.T) { + s := &HumanReadableSerializer{} + input := hrEdgeIgnoreField{A: "shown", B: "hidden", C: "default"} + + data, err := s.Marshal(input) + require.NoError(t, err) + + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + assert.Equal(t, "shown", raw["a"]) + _, hasB := raw["B"] + assert.False(t, hasB, `json:"-" field must not be serialized`) + assert.Equal(t, "default", raw["C"]) + + var got hrEdgeIgnoreField + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, "shown", got.A) + assert.Equal(t, "", got.B, `json:"-" field must remain zero on decode`) + assert.Equal(t, "default", got.C) +} + +func TestHumanReadableSerializer_StructFieldFallbackToFieldName(t *testing.T) { + s := &HumanReadableSerializer{} + + // Hand-craft JSON using the Go field name (no tag). hrUnmarshalStruct should + // look up `data[fieldName]` first, then fall back to `data[field.Name]`. + raw := []byte(`{"Name":"x","value":99}`) + var got hrTestStruct + require.NoError(t, s.Unmarshal(raw, &got)) + assert.Equal(t, "x", got.Name) + assert.Equal(t, 99, got.Value) +} + +func TestHumanReadableSerializer_ConcreteFieldDoesNotRequireRegistration(t *testing.T) { + _ = GenericRegister[hrEdgeUnregisteredField]("hr_edge_unregistered_field") + + s := &HumanReadableSerializer{} + input := hrEdgeUnregisteredField{V: hrUnregisteredInner{N: 7}} + + data, err := s.Marshal(input) + require.NoError(t, err, "concrete struct field shouldn't require its element type to be registered") + + var got hrEdgeUnregisteredField + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input, got) +} + +// ============================================================================= +// Section: Error paths +// ============================================================================= + +func TestHumanReadableSerializer_UnmarshalErrors(t *testing.T) { + s := &HumanReadableSerializer{} + + t.Run("corrupt JSON", func(t *testing.T) { + var got hrTestStruct + err := s.Unmarshal([]byte(`{"name":`), &got) + require.Error(t, err) + assert.Contains(t, err.Error(), "unmarshal JSON") + }) + + t.Run("nil pointer target", func(t *testing.T) { + var ptr *hrTestStruct + err := s.Unmarshal([]byte(`{}`), ptr) + require.Error(t, err) + assert.Contains(t, err.Error(), "non-nil pointer") + }) + + t.Run("non-pointer target", func(t *testing.T) { + var v hrTestStruct + err := s.Unmarshal([]byte(`{}`), v) + require.Error(t, err) + assert.Contains(t, err.Error(), "non-nil pointer") + }) + + t.Run("unknown $type", func(t *testing.T) { + var got hrEdgeAnyContainer + err := s.Unmarshal([]byte(`{"v":{"$type":"this_type_is_not_registered","value":1}}`), &got) + // shouldTreatAsTypeEnvelope returns false for unknown type names, so the + // payload is passed through as a plain map[string]any. + require.NoError(t, err) + m, ok := got.V.(map[string]any) + require.True(t, ok) + assert.Equal(t, "this_type_is_not_registered", m["$type"]) + }) + + t.Run("typed envelope with bad inner data", func(t *testing.T) { + var got hrEdgeAnyContainer + // `_eino_int` expects a numeric value; a JSON object cannot decode into int. + err := s.Unmarshal([]byte(`{"v":{"$type":"_eino_int","value":{"oops":1}}}`), &got) + require.Error(t, err) + }) + + t.Run("array on a non-slice/array target", func(t *testing.T) { + var got hrTestStruct + err := s.Unmarshal([]byte(`[1,2,3]`), &got) + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot unmarshal slice") + }) + + t.Run("object on a non-map/struct target", func(t *testing.T) { + var got int + err := s.Unmarshal([]byte(`{"a":1}`), &got) + require.Error(t, err) + }) +} + +func TestHumanReadableSerializer_MarshalErrors(t *testing.T) { + s := &HumanReadableSerializer{} + + t.Run("unregistered type via interface field", func(t *testing.T) { + input := hrEdgeAnyContainer{V: hrUnregisteredHere{X: 1}} + _, err := s.Marshal(input) + require.Error(t, err) + assert.Contains(t, err.Error(), "unknown type") + }) + + t.Run("array of unregistered element via interface field", func(t *testing.T) { + input := hrEdgeAnyContainer{V: [2]hrUnregisteredHere{{X: 1}, {X: 2}}} + _, err := s.Marshal(input) + require.Error(t, err) + }) +} + +// ============================================================================= +// Section: Top-level primitives +// ============================================================================= + +func TestHumanReadableSerializer_TopLevelPrimitives(t *testing.T) { + s := &HumanReadableSerializer{} + + t.Run("int", func(t *testing.T) { + data, err := s.Marshal(int(42)) + require.NoError(t, err) + var got int + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, 42, got) + }) + + t.Run("float", func(t *testing.T) { + data, err := s.Marshal(3.14) + require.NoError(t, err) + var got float64 + require.NoError(t, s.Unmarshal(data, &got)) + assert.InDelta(t, 3.14, got, 1e-9) + }) + + t.Run("string", func(t *testing.T) { + data, err := s.Marshal("hello") + require.NoError(t, err) + var got string + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, "hello", got) + }) + + t.Run("[]int", func(t *testing.T) { + data, err := s.Marshal([]int{1, 2, 3}) + require.NoError(t, err) + var got []int + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, []int{1, 2, 3}, got) + }) + + t.Run("map[string]int", func(t *testing.T) { + data, err := s.Marshal(map[string]int{"a": 1, "b": 2}) + require.NoError(t, err) + var got map[string]int + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, map[string]int{"a": 1, "b": 2}, got) + }) +} + +// ============================================================================= +// Section: Internal helpers +// ============================================================================= + +func TestIsEmptyValue_AllKinds(t *testing.T) { + cases := []struct { + name string + v any + want bool + }{ + {"empty string", "", true}, + {"non-empty string", "x", false}, + {"empty slice", []int{}, true}, + {"non-empty slice", []int{1}, false}, + {"nil slice", []int(nil), true}, + {"empty map", map[string]int{}, true}, + {"non-empty map", map[string]int{"a": 1}, false}, + {"empty array", [0]int{}, true}, + {"non-empty array", [3]int{1, 2, 3}, false}, + {"false bool", false, true}, + {"true bool", true, false}, + {"int 0", int(0), true}, + {"int non-zero", int(5), false}, + {"int8 0", int8(0), true}, + {"uint 0", uint(0), true}, + {"uint64 0", uint64(0), true}, + {"float64 0", float64(0), true}, + {"float64 non-zero", float64(0.5), false}, + {"nil pointer", (*int)(nil), true}, + {"non-nil pointer", func() any { v := 1; return &v }(), false}, + {"nil interface", any(nil), true}, + // Channel hits the `default: return false` branch. + {"channel (default branch)", make(chan int), false}, + {"func (default branch)", func() {}, false}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + rv := reflect.ValueOf(c.v) + if !rv.IsValid() { + assert.Equal(t, c.want, true) + return + } + got := isEmptyValue(rv) + assert.Equal(t, c.want, got) + }) + } +} + +func TestSetValueWithConversion_AllPaths(t *testing.T) { + t.Run("invalid source sets zero", func(t *testing.T) { + var dst int = 99 + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.Value{}) + assert.True(t, ok) + assert.Equal(t, 0, dst, "invalid source must zero the target") + }) + + t.Run("ptr target nil → allocates", func(t *testing.T) { + var p *int + target := reflect.ValueOf(&p).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(42)) + assert.True(t, ok) + require.NotNil(t, p) + assert.Equal(t, 42, *p) + }) + + t.Run("ptr source nil → target zeroed", func(t *testing.T) { + var src *int + var dst int = 99 + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(src)) + assert.True(t, ok) + assert.Equal(t, 0, dst) + }) + + t.Run("ptr source non-nil → deref then set", func(t *testing.T) { + v := 7 + var dst int + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(&v)) + assert.True(t, ok) + assert.Equal(t, 7, dst) + }) + + t.Run("convertible types", func(t *testing.T) { + var dst int32 + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(int64(100))) + assert.True(t, ok) + assert.Equal(t, int32(100), dst) + }) + + t.Run("float64 → int", func(t *testing.T) { + var dst int + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(float64(7.0))) + assert.True(t, ok) + assert.Equal(t, 7, dst) + }) + + t.Run("int → int (different bit widths)", func(t *testing.T) { + var dst int64 + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(int(42))) + assert.True(t, ok) + assert.Equal(t, int64(42), dst) + }) + + t.Run("float64 → uint", func(t *testing.T) { + var dst uint + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(float64(8))) + assert.True(t, ok) + assert.Equal(t, uint(8), dst) + }) + + t.Run("int → float", func(t *testing.T) { + var dst float64 + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf(int(12))) + assert.True(t, ok) + assert.Equal(t, float64(12), dst) + }) + + t.Run("incompatible types return false", func(t *testing.T) { + var dst struct{ A int } + target := reflect.ValueOf(&dst).Elem() + ok := setValueWithConversion(target, reflect.ValueOf("not a struct")) + assert.False(t, ok) + }) +} + +func TestGetJSONFieldName_Variants(t *testing.T) { + assert.Equal(t, "Name", getJSONFieldName("Name", "")) + assert.Equal(t, "alias", getJSONFieldName("Name", "alias")) + assert.Equal(t, "alias", getJSONFieldName("Name", "alias,omitempty")) + // Empty primary part (only ",omitempty") falls back to field name. + assert.Equal(t, "Name", getJSONFieldName("Name", ",omitempty")) +} + +func TestParseTypeName_Errors(t *testing.T) { + t.Run("unknown plain type", func(t *testing.T) { + _, _, err := parseTypeName("not_registered") + require.Error(t, err) + }) + + t.Run("malformed array missing close bracket", func(t *testing.T) { + _, _, err := parseTypeName("[3 _eino_int") + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid array") + }) + + t.Run("array with bad size", func(t *testing.T) { + _, _, err := parseTypeName("[abc]_eino_int") + require.Error(t, err) + }) + + t.Run("array with unknown elem", func(t *testing.T) { + _, _, err := parseTypeName("[3]not_registered") + require.Error(t, err) + }) + + t.Run("slice with unknown elem", func(t *testing.T) { + _, _, err := parseTypeName("[]not_registered") + require.Error(t, err) + }) + + t.Run("map with unknown key type", func(t *testing.T) { + _, _, err := parseTypeName("map[not_registered]_eino_string") + require.Error(t, err) + assert.Contains(t, err.Error(), "key") + }) + + t.Run("map with unknown value type", func(t *testing.T) { + _, _, err := parseTypeName("map[_eino_string]not_registered") + require.Error(t, err) + assert.Contains(t, err.Error(), "value") + }) + + t.Run("nested map with pointer key/value", func(t *testing.T) { + rt, ptr, err := parseTypeName("map[*_eino_string]*_eino_int") + require.NoError(t, err) + assert.Equal(t, uint32(0), ptr) + assert.Equal(t, reflect.Map, rt.Kind()) + assert.Equal(t, reflect.Ptr, rt.Key().Kind()) + assert.Equal(t, reflect.Ptr, rt.Elem().Kind()) + }) + + t.Run("pointer to slice", func(t *testing.T) { + rt, ptr, err := parseTypeName("*[]_eino_int") + require.NoError(t, err) + assert.Equal(t, uint32(1), ptr) + assert.Equal(t, reflect.Slice, rt.Kind()) + }) +} + +func TestGetTypeName_AllShapes(t *testing.T) { + // Plain registered. + n, err := getTypeName(reflect.TypeOf(int(0))) + require.NoError(t, err) + assert.Equal(t, "_eino_int", n) + + // Pointer. + n, err = getTypeName(reflect.TypeOf((*int)(nil))) + require.NoError(t, err) + assert.Equal(t, "*_eino_int", n) + + // Slice. + n, err = getTypeName(reflect.TypeOf([]int{})) + require.NoError(t, err) + assert.Equal(t, "[]_eino_int", n) + + // Array. + n, err = getTypeName(reflect.TypeOf([3]int{})) + require.NoError(t, err) + assert.Equal(t, "[3]_eino_int", n) + + // Map. + n, err = getTypeName(reflect.TypeOf(map[string]int{})) + require.NoError(t, err) + assert.Equal(t, "map[_eino_string]_eino_int", n) + + // Unregistered. + type unreg struct{} + _, err = getTypeName(reflect.TypeOf(unreg{})) + require.Error(t, err) + + // Slice with unregistered elem. + _, err = getTypeName(reflect.TypeOf([]unreg{})) + require.Error(t, err) + + // Array with unregistered elem. + _, err = getTypeName(reflect.TypeOf([3]unreg{})) + require.Error(t, err) + + // Map with unregistered key. + _, err = getTypeName(reflect.TypeOf(map[unreg]int{})) + require.Error(t, err) + + // Map with unregistered value. + _, err = getTypeName(reflect.TypeOf(map[string]unreg{})) + require.Error(t, err) +} + +// ============================================================================= +// Section: Protocol safety +// ============================================================================= + +func TestHumanReadableSerializer_NilSliceField(t *testing.T) { + type holder struct { + S []int `json:"s"` + } + _ = GenericRegister[holder]("hr_edge_nil_slice_holder") + + s := &HumanReadableSerializer{} + input := holder{S: nil} + + data, err := s.Marshal(input) + require.NoError(t, err) + + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + assert.Nil(t, raw["s"], "nil slice must serialize as JSON null") + + var got holder + require.NoError(t, s.Unmarshal(data, &got)) + assert.Nil(t, got.S, "JSON null must decode back to nil slice") +} + +func TestHumanReadableSerializer_NilMapField(t *testing.T) { + type holder struct { + M map[string]int `json:"m"` + } + _ = GenericRegister[holder]("hr_edge_nil_map_holder") + + s := &HumanReadableSerializer{} + input := holder{M: nil} + + data, err := s.Marshal(input) + require.NoError(t, err) + + var got holder + require.NoError(t, s.Unmarshal(data, &got)) + assert.Nil(t, got.M) +} + +func TestHumanReadableSerializer_MapWithUnregisteredValueInInterface(t *testing.T) { + type unregValue struct{ N int } + type holder struct { + V any `json:"v"` + } + _ = GenericRegister[holder]("hr_edge_unreg_map_value_holder") + + s := &HumanReadableSerializer{} + input := holder{V: map[string]unregValue{"a": {N: 1}}} + + _, err := s.Marshal(input) + require.Error(t, err, "map with unregistered value type in interface field must error") +} + +func TestHumanReadableSerializer_StringWithSpecialCharacters(t *testing.T) { + type holder struct { + S string `json:"s"` + A any `json:"a"` + } + _ = GenericRegister[holder]("hr_edge_special_chars_holder") + + cases := []string{ + `"quotes"`, + `back\slash`, + "newline\nand\ttab", + "unicode 你好 🚀", + "control\x01\x02\x03", + "", + "$type:should-not-confuse-parser", + } + s := &HumanReadableSerializer{} + for _, c := range cases { + t.Run(c, func(t *testing.T) { + input := holder{S: c, A: c} + data, err := s.Marshal(input) + require.NoError(t, err) + var got holder + require.NoError(t, s.Unmarshal(data, &got)) + assert.Equal(t, input.S, got.S) + assert.Equal(t, input.A, got.A) + }) + } +} + +func TestHumanReadableSerializer_DeepRecursion(t *testing.T) { + type node struct { + V int `json:"v"` + Next *node `json:"next,omitempty"` + } + _ = GenericRegister[node]("hr_edge_deep_node") + + s := &HumanReadableSerializer{} + + // Build a chain of 50 nodes. + const depth = 50 + root := &node{V: 0} + cur := root + for i := 1; i < depth; i++ { + cur.Next = &node{V: i} + cur = cur.Next + } + + data, err := s.Marshal(root) + require.NoError(t, err) + + var got node + require.NoError(t, s.Unmarshal(data, &got)) + + // Walk and verify all values. + cur = &got + for i := 0; i < depth; i++ { + require.NotNil(t, cur, "node at depth %d", i) + assert.Equal(t, i, cur.V) + cur = cur.Next + } +} diff --git a/internal/serialization/serialization_benchmark_test.go b/internal/serialization/serialization_benchmark_test.go new file mode 100644 index 000000000..623854d87 --- /dev/null +++ b/internal/serialization/serialization_benchmark_test.go @@ -0,0 +1,613 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package serialization + +import ( + "encoding/gob" + "fmt" + "testing" +) + +type benchMessage struct { + Role string `json:"role"` + Content string `json:"content"` + Name string `json:"name,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + ReasoningContent string `json:"reasoning_content,omitempty"` + Extra map[string]any `json:"extra,omitempty"` +} + +type benchToolCall struct { + ID string `json:"id"` + Type string `json:"type"` + Function benchFunctionCall `json:"function"` + Extra map[string]any `json:"extra,omitempty"` +} + +type benchFunctionCall struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +type benchCustomType struct { + Provider string `json:"provider"` + Model string `json:"model"` + Version int `json:"version"` +} + +type benchStructWithInterface struct { + A any + B any + C map[string]any +} + +func init() { + _ = GenericRegister[benchMessage]("bench_message") + _ = GenericRegister[benchToolCall]("bench_tool_call") + _ = GenericRegister[benchFunctionCall]("bench_function_call") + _ = GenericRegister[benchCustomType]("bench_custom_type") + _ = GenericRegister[benchStructWithInterface]("bench_struct_with_interface") + + gob.Register(benchMessage{}) + gob.Register(benchToolCall{}) + gob.Register(benchFunctionCall{}) + gob.Register(benchCustomType{}) + gob.Register(benchStructWithInterface{}) +} + +func createSimpleMessage() benchMessage { + return benchMessage{ + Role: "assistant", + Content: "Hello, how can I help you today?", + Extra: map[string]any{ + "model": "gpt-4", + "temperature": 0.7, + "max_tokens": 1024, + }, + } +} + +func createComplexMessage() benchMessage { + return benchMessage{ + Role: "assistant", + Content: "Here is a detailed response with multiple paragraphs of content that simulates a real-world LLM response. This includes various information and explanations that would typically be returned by a language model.", + Name: "assistant", + ReasoningContent: "Let me think about this step by step. First, I need to understand the question. Then I'll formulate a comprehensive response.", + Extra: map[string]any{ + "model": "gpt-4-turbo", + "temperature": 0.7, + "max_tokens": 4096, + "top_p": 0.95, + "frequency_penalty": 0.0, + "presence_penalty": 0.0, + "stop_sequences": []string{"END", "STOP"}, + "metadata": benchCustomType{ + Provider: "benchmark", + Model: "metadata", + Version: 1, + }, + }, + } +} + +func createMessageWithCustomTypes() benchStructWithInterface { + return benchStructWithInterface{ + A: "simple string", + B: benchCustomType{ + Provider: "openai", + Model: "gpt-4", + Version: 4, + }, + C: map[string]any{ + "config": benchCustomType{ + Provider: "anthropic", + Model: "claude-3", + Version: 3, + }, + "count": 42, + }, + } +} + +func createLargeMessage() benchMessage { + extra := make(map[string]any) + for i := 0; i < 50; i++ { + extra[fmt.Sprintf("key_%d", i)] = fmt.Sprintf("value_%d", i) + } + extra["nested"] = benchCustomType{Provider: "benchmark", Model: "nested", Version: 3} + extra["list"] = []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} + + return benchMessage{ + Role: "assistant", + Content: "This is a large message with many extra fields to test serialization performance with larger payloads.", + Extra: extra, + } +} + +func BenchmarkInternalSerializer_Marshal_SimpleMessage(b *testing.B) { + s := &InternalSerializer{} + msg := createSimpleMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Marshal_SimpleMessage(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createSimpleMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_Unmarshal_SimpleMessage(b *testing.B) { + s := &InternalSerializer{} + msg := createSimpleMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Unmarshal_SimpleMessage(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createSimpleMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_Marshal_ComplexMessage(b *testing.B) { + s := &InternalSerializer{} + msg := createComplexMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Marshal_ComplexMessage(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createComplexMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_Unmarshal_ComplexMessage(b *testing.B) { + s := &InternalSerializer{} + msg := createComplexMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Unmarshal_ComplexMessage(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createComplexMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_Marshal_CustomTypes(b *testing.B) { + s := &InternalSerializer{} + msg := createMessageWithCustomTypes() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Marshal_CustomTypes(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createMessageWithCustomTypes() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_Unmarshal_CustomTypes(b *testing.B) { + s := &InternalSerializer{} + msg := createMessageWithCustomTypes() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchStructWithInterface + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Unmarshal_CustomTypes(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createMessageWithCustomTypes() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchStructWithInterface + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_Marshal_LargeMessage(b *testing.B) { + s := &InternalSerializer{} + msg := createLargeMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Marshal_LargeMessage(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createLargeMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_Unmarshal_LargeMessage(b *testing.B) { + s := &InternalSerializer{} + msg := createLargeMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_Unmarshal_LargeMessage(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createLargeMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkInternalSerializer_RoundTrip_SimpleMessage(b *testing.B) { + s := &InternalSerializer{} + msg := createSimpleMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + data, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + var result benchMessage + err = s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkHumanReadableSerializer_RoundTrip_SimpleMessage(b *testing.B) { + s := &HumanReadableSerializer{} + msg := createSimpleMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + data, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + var result benchMessage + err = s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Marshal_SimpleMessage(b *testing.B) { + s := &GobSerializer{} + msg := createSimpleMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Unmarshal_SimpleMessage(b *testing.B) { + s := &GobSerializer{} + msg := createSimpleMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Marshal_ComplexMessage(b *testing.B) { + s := &GobSerializer{} + msg := createComplexMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Unmarshal_ComplexMessage(b *testing.B) { + s := &GobSerializer{} + msg := createComplexMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Marshal_CustomTypes(b *testing.B) { + s := &GobSerializer{} + msg := createMessageWithCustomTypes() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Unmarshal_CustomTypes(b *testing.B) { + s := &GobSerializer{} + msg := createMessageWithCustomTypes() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchStructWithInterface + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Marshal_LargeMessage(b *testing.B) { + s := &GobSerializer{} + msg := createLargeMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + _, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_Unmarshal_LargeMessage(b *testing.B) { + s := &GobSerializer{} + msg := createLargeMessage() + data, _ := s.Marshal(msg) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + var result benchMessage + err := s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkGobSerializer_RoundTrip_SimpleMessage(b *testing.B) { + s := &GobSerializer{} + msg := createSimpleMessage() + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + data, err := s.Marshal(msg) + if err != nil { + b.Fatal(err) + } + var result benchMessage + err = s.Unmarshal(data, &result) + if err != nil { + b.Fatal(err) + } + } +} + +func TestOutputSizeComparison(t *testing.T) { + is := &InternalSerializer{} + hr := &HumanReadableSerializer{} + gs := &GobSerializer{} + + testCases := []struct { + name string + input any + }{ + {"SimpleMessage", createSimpleMessage()}, + {"ComplexMessage", createComplexMessage()}, + {"CustomTypes", createMessageWithCustomTypes()}, + {"LargeMessage", createLargeMessage()}, + } + + for _, tc := range testCases { + isData, _ := is.Marshal(tc.input) + hrData, _ := hr.Marshal(tc.input) + gsData, _ := gs.Marshal(tc.input) + + t.Logf("%s:", tc.name) + t.Logf(" InternalSerializer: %d bytes", len(isData)) + t.Logf(" HumanReadableSerializer: %d bytes", len(hrData)) + t.Logf(" GobSerializer: %d bytes", len(gsData)) + t.Logf(" Ratio (HR/IS): %.2f%%", float64(len(hrData))/float64(len(isData))*100) + t.Logf(" Ratio (Gob/IS): %.2f%%", float64(len(gsData))/float64(len(isData))*100) + t.Logf(" HumanReadable output:\n%s\n", string(hrData)) + } +} diff --git a/internal/serialization/serialization_test.go b/internal/serialization/serialization_test.go index 014c7fd94..e9fae398c 100644 --- a/internal/serialization/serialization_test.go +++ b/internal/serialization/serialization_test.go @@ -24,6 +24,18 @@ import ( "github.com/stretchr/testify/require" ) +type Serializer interface { + Marshal(v any) ([]byte, error) + Unmarshal(data []byte, v any) error +} + +func getSerializers() map[string]Serializer { + return map[string]Serializer{ + "InternalSerializer": &InternalSerializer{}, + "HumanReadableSerializer": &HumanReadableSerializer{}, + } +} + type myInterface interface { Method() } @@ -66,10 +78,35 @@ func (m myStruct4) MarshalJSON() ([]byte, error) { return []byte(m.FieldA), nil } -func TestSerialization(t *testing.T) { +type myStruct5 struct { + FieldA string +} + +func (m *myStruct5) UnmarshalJSON(bytes []byte) error { + m.FieldA = "FieldA" + return nil +} + +func (m myStruct5) MarshalJSON() ([]byte, error) { + return []byte("1"), nil +} + +type unmarshalTestStruct struct { + Foo string + Bar int +} + +func init() { _ = GenericRegister[myStruct]("myStruct") _ = GenericRegister[myStruct2]("myStruct2") + _ = GenericRegister[myStruct3]("myStruct3") + _ = GenericRegister[myStruct4]("myStruct4") + _ = GenericRegister[myStruct5]("myStruct5") _ = GenericRegister[myInterface]("myInterface") + _ = GenericRegister[unmarshalTestStruct]("unmarshalTestStruct") +} + +func TestSerialization_RoundTrip(t *testing.T) { ms := myStruct{A: "test"} pms := &ms pointerOfPointerOfMyStruct := &pms @@ -78,274 +115,507 @@ func TestSerialization(t *testing.T) { ms2 := myStruct{A: "2"} ms3 := myStruct{A: "3"} ms4 := myStruct{A: "4"} - values := []any{ - 10, - "test", - ms, - pms, - pointerOfPointerOfMyStruct, - myInterface(pms), - []int{1, 2, 3}, - []any{1, "test"}, - []myInterface{nil, &myStruct{A: "1"}, &myStruct{A: "2"}}, - map[string]string{"123": "123", "abc": "abc"}, - map[string]myInterface{"1": nil, "2": pms}, - map[string]any{"123": 1, "abc": &myStruct{A: "1"}, "bcd": nil}, - map[myStruct]any{ + + testCases := []struct { + name string + value any + }{ + {"int", 10}, + {"string", "test"}, + {"struct", ms}, + {"pointer to struct", pms}, + {"pointer to pointer of struct", pointerOfPointerOfMyStruct}, + {"interface", myInterface(pms)}, + {"slice of int", []int{1, 2, 3}}, + {"slice of any", []any{1, "test"}}, + {"slice of interface with nil", []myInterface{nil, &myStruct{A: "1"}, &myStruct{A: "2"}}}, + {"map string to string", map[string]string{"123": "123", "abc": "abc"}}, + {"map string to interface with nil", map[string]myInterface{"1": nil, "2": pms}}, + {"map string to any with nil", map[string]any{"123": 1, "abc": &myStruct{A: "1"}, "bcd": nil}}, + {"map struct to any complex", map[myStruct]any{ ms1: 1, - ms2: &myStruct{ - A: "2", - }, + ms2: &myStruct{A: "2"}, ms3: nil, ms4: []any{ 1, pointerOfPointerOfMyStruct, - "123", &myStruct{ - A: "1", - }, + "123", + &myStruct{A: "1"}, nil, map[myStruct]any{ ms1: 1, ms2: nil, }, }, - }, - myStruct2{ + }}, + {"complex struct", myStruct2{ A: "123", - B: &myStruct{ - A: "test", - }, - C: map[string]**myStruct{ - "a": pointerOfPointerOfMyStruct, - }, + B: &myStruct{A: "test"}, + C: map[string]**myStruct{"a": pointerOfPointerOfMyStruct}, D: map[myStruct]any{{"a"}: 1}, E: []any{1, "2", 3}, f: "", - G: myStruct3{ - FieldA: "1", - }, + G: myStruct3{FieldA: "1"}, H: nil, - I: []*myStruct3{ - {FieldA: "2"}, {FieldA: "3"}, - }, - J: map[string]myStruct3{ - "1": {FieldA: "4"}, - "2": {FieldA: "5"}, - }, - K: myStruct4{ - FieldA: "1", - }, - L: []*myStruct4{ - {FieldA: "2"}, {FieldA: "3"}, - }, - M: map[string]myStruct4{ - "1": {FieldA: "4"}, - "2": {FieldA: "5"}, - }, - }, - map[string]map[string][]map[string][][]string{ + I: []*myStruct3{{FieldA: "2"}, {FieldA: "3"}}, + J: map[string]myStruct3{"1": {FieldA: "4"}, "2": {FieldA: "5"}}, + K: myStruct4{FieldA: "1"}, + L: []*myStruct4{{FieldA: "2"}, {FieldA: "3"}}, + M: map[string]myStruct4{"1": {FieldA: "4"}, "2": {FieldA: "5"}}, + }}, + {"deeply nested map", map[string]map[string][]map[string][][]string{ "1": { "a": []map[string][][]string{ - {"b": { - {"c"}, - {"d"}, - }}, + {"b": {{"c"}, {"d"}}}, }, }, - }, - []*myStruct{}, - &myStruct{}, + }}, + {"empty slice of pointers", []*myStruct{}}, + {"empty struct pointer", &myStruct{}}, } - for _, value := range values { - data, err := (&InternalSerializer{}).Marshal(value) - assert.NoError(t, err) - v := reflect.New(reflect.TypeOf(value)).Interface() - err = (&InternalSerializer{}).Unmarshal(data, v) - assert.NoError(t, err) - assert.Equal(t, value, reflect.ValueOf(v).Elem().Interface()) + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + data, err := s.Marshal(tc.value) + require.NoError(t, err, "marshal failed") + + result := reflect.New(reflect.TypeOf(tc.value)).Interface() + err = s.Unmarshal(data, result) + require.NoError(t, err, "unmarshal failed") + + assert.Equal(t, tc.value, reflect.ValueOf(result).Elem().Interface()) + }) + } + }) } } -type myStruct5 struct { - FieldA string -} +func TestSerialization_CustomMarshaler(t *testing.T) { + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + t.Run("struct with custom marshaler", func(t *testing.T) { + input := myStruct5{FieldA: "1"} + data, err := s.Marshal(input) + require.NoError(t, err) -func (m *myStruct5) UnmarshalJSON(bytes []byte) error { - m.FieldA = "FieldA" - return nil + result := &myStruct5{} + err = s.Unmarshal(data, result) + require.NoError(t, err) + assert.Equal(t, myStruct5{FieldA: "FieldA"}, *result) + }) + + t.Run("custom marshaler in map[string]any", func(t *testing.T) { + input := map[string]any{ + "1": myStruct5{FieldA: "1"}, + } + data, err := s.Marshal(input) + require.NoError(t, err) + + result := map[string]any{} + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Equal(t, map[string]any{ + "1": myStruct5{FieldA: "FieldA"}, + }, result) + }) + }) + } } -func (m myStruct5) MarshalJSON() ([]byte, error) { - return []byte("1"), nil +func TestSerialization_Unmarshal(t *testing.T) { + ptr := func(i int) *int { return &i } + + successCases := []struct { + name string + inputValue any + outputPtr func() any + expectedVal any + }{ + { + name: "simple type", + inputValue: 123, + outputPtr: func() any { return new(int) }, + expectedVal: 123, + }, + { + name: "struct type", + inputValue: unmarshalTestStruct{Foo: "hello", Bar: 42}, + outputPtr: func() any { return new(unmarshalTestStruct) }, + expectedVal: unmarshalTestStruct{Foo: "hello", Bar: 42}, + }, + { + name: "pointer to struct", + inputValue: &unmarshalTestStruct{Foo: "world", Bar: 99}, + outputPtr: func() any { return new(*unmarshalTestStruct) }, + expectedVal: &unmarshalTestStruct{Foo: "world", Bar: 99}, + }, + { + name: "unmarshal pointer to value", + inputValue: &unmarshalTestStruct{Foo: "p2v", Bar: 1}, + outputPtr: func() any { return new(unmarshalTestStruct) }, + expectedVal: unmarshalTestStruct{Foo: "p2v", Bar: 1}, + }, + { + name: "unmarshal value to pointer", + inputValue: unmarshalTestStruct{Foo: "v2p", Bar: 2}, + outputPtr: func() any { return new(*unmarshalTestStruct) }, + expectedVal: &unmarshalTestStruct{Foo: "v2p", Bar: 2}, + }, + { + name: "convertible types", + inputValue: int32(42), + outputPtr: func() any { return new(int64) }, + expectedVal: int64(42), + }, + { + name: "pointer to pointer destination", + inputValue: 12345, + outputPtr: func() any { return new(*int) }, + expectedVal: ptr(12345), + }, + { + name: "unmarshal to any", + inputValue: unmarshalTestStruct{Foo: "any", Bar: 101}, + outputPtr: func() any { return new(any) }, + expectedVal: unmarshalTestStruct{Foo: "any", Bar: 101}, + }, + } + + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + t.Run("success cases", func(t *testing.T) { + for _, tc := range successCases { + t.Run(tc.name, func(t *testing.T) { + data, err := s.Marshal(tc.inputValue) + require.NoError(t, err) + + outputPtr := tc.outputPtr() + err = s.Unmarshal(data, outputPtr) + require.NoError(t, err) + + actualVal := reflect.ValueOf(outputPtr).Elem().Interface() + assert.Equal(t, tc.expectedVal, actualVal) + }) + } + }) + + t.Run("unmarshal nil pointer", func(t *testing.T) { + data, err := s.Marshal((*unmarshalTestStruct)(nil)) + require.NoError(t, err) + + var result *unmarshalTestStruct = &unmarshalTestStruct{} + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Nil(t, result) + }) + + t.Run("error cases", func(t *testing.T) { + data, err := s.Marshal(123) + require.NoError(t, err) + + t.Run("destination not a pointer", func(t *testing.T) { + var output int + err := s.Unmarshal(data, output) + require.Error(t, err) + assert.Contains(t, err.Error(), "non-nil pointer") + }) + + t.Run("destination is a nil pointer", func(t *testing.T) { + var output *int + err := s.Unmarshal(data, output) + require.Error(t, err) + assert.Contains(t, err.Error(), "non-nil pointer") + }) + + t.Run("type mismatch", func(t *testing.T) { + strData, mErr := s.Marshal("i am a string") + require.NoError(t, mErr) + + var output int + err := s.Unmarshal(strData, &output) + require.Error(t, err) + }) + + t.Run("unconvertible types", func(t *testing.T) { + intData, mErr := s.Marshal(123) + require.NoError(t, mErr) + + var output bool + err := s.Unmarshal(intData, &output) + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot assign") + }) + }) + }) + } } -func TestMarshalStruct(t *testing.T) { - assert.NoError(t, GenericRegister[myStruct5]("myStruct5")) - s := myStruct5{FieldA: "1"} - data, err := (&InternalSerializer{}).Marshal(s) - assert.NoError(t, err) - result := &myStruct5{} - err = (&InternalSerializer{}).Unmarshal(data, result) - assert.NoError(t, err) - assert.Equal(t, myStruct5{FieldA: "FieldA"}, *result) - - ma := map[string]any{ - "1": s, +func TestSerialization_PrimitiveTypes(t *testing.T) { + testCases := []struct { + name string + value any + }{ + {"int", int(42)}, + {"int8", int8(8)}, + {"int16", int16(16)}, + {"int32", int32(32)}, + {"int64", int64(64)}, + {"uint", uint(42)}, + {"uint8", uint8(8)}, + {"uint16", uint16(16)}, + {"uint32", uint32(32)}, + {"uint64", uint64(64)}, + {"float32", float32(3.14)}, + {"float64", float64(3.14159)}, + {"bool true", true}, + {"bool false", false}, + {"string", "hello world"}, + {"empty string", ""}, + } + + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + data, err := s.Marshal(tc.value) + require.NoError(t, err) + + result := reflect.New(reflect.TypeOf(tc.value)).Interface() + err = s.Unmarshal(data, result) + require.NoError(t, err) + + assert.Equal(t, tc.value, reflect.ValueOf(result).Elem().Interface()) + }) + } + }) } - data, err = (&InternalSerializer{}).Marshal(ma) - assert.NoError(t, err) - result2 := map[string]any{} - err = (&InternalSerializer{}).Unmarshal(data, &result2) - assert.NoError(t, err) - assert.Equal(t, map[string]any{ - "1": myStruct5{FieldA: "FieldA"}, - }, result2) } -type unmarshalTestStruct struct { - Foo string - Bar int +func TestSerialization_NilValues(t *testing.T) { + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + t.Run("nil pointer", func(t *testing.T) { + var input *myStruct = nil + data, err := s.Marshal(input) + require.NoError(t, err) + + var result *myStruct + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Nil(t, result) + }) + + t.Run("nil in slice", func(t *testing.T) { + input := []any{1, nil, "three", nil, 5} + data, err := s.Marshal(input) + require.NoError(t, err) + + var result []any + err = s.Unmarshal(data, &result) + require.NoError(t, err) + require.Len(t, result, 5) + assert.Equal(t, 1, result[0]) + assert.Nil(t, result[1]) + assert.Equal(t, "three", result[2]) + assert.Nil(t, result[3]) + assert.Equal(t, 5, result[4]) + }) + + t.Run("nil in map", func(t *testing.T) { + input := map[string]any{ + "value": 123, + "nil": nil, + } + data, err := s.Marshal(input) + require.NoError(t, err) + + var result map[string]any + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Equal(t, 123, result["value"]) + assert.Nil(t, result["nil"]) + }) + }) + } } -func init() { - // Register types for the serializer to work. - // This is necessary for the serializer to know how to handle custom struct types. - err := GenericRegister[unmarshalTestStruct]("unmarshalTestStruct") - if err != nil { - panic(err) +func TestSerialization_EmptyCollections(t *testing.T) { + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + t.Run("empty slice", func(t *testing.T) { + input := []int{} + data, err := s.Marshal(input) + require.NoError(t, err) + + var result []int + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.NotNil(t, result) + assert.Len(t, result, 0) + }) + + t.Run("empty map", func(t *testing.T) { + input := map[string]any{} + data, err := s.Marshal(input) + require.NoError(t, err) + + var result map[string]any + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.NotNil(t, result) + assert.Len(t, result, 0) + }) + + t.Run("empty slice of pointers", func(t *testing.T) { + input := []*myStruct{} + data, err := s.Marshal(input) + require.NoError(t, err) + + var result []*myStruct + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.NotNil(t, result) + assert.Len(t, result, 0) + }) + }) } } -func TestInternalSerializer_Unmarshal(t *testing.T) { - s := InternalSerializer{} - - t.Run("success cases", func(t *testing.T) { - // Helper to create a pointer to a value, needed for the expected value in one test case. - ptr := func(i int) *int { return &i } - - testCases := []struct { - name string - inputValue any - outputPtr any - expectedVal any - }{ - { - name: "simple type", - inputValue: 123, - outputPtr: new(int), - expectedVal: 123, - }, - { - name: "struct type", - inputValue: unmarshalTestStruct{Foo: "hello", Bar: 42}, - outputPtr: new(unmarshalTestStruct), - expectedVal: unmarshalTestStruct{Foo: "hello", Bar: 42}, - }, - { - name: "pointer to struct", - inputValue: &unmarshalTestStruct{Foo: "world", Bar: 99}, - outputPtr: new(*unmarshalTestStruct), - expectedVal: &unmarshalTestStruct{Foo: "world", Bar: 99}, - }, - { - name: "unmarshal pointer to value", - inputValue: &unmarshalTestStruct{Foo: "p2v", Bar: 1}, - outputPtr: new(unmarshalTestStruct), - expectedVal: unmarshalTestStruct{Foo: "p2v", Bar: 1}, - }, - { - name: "unmarshal value to pointer", - inputValue: unmarshalTestStruct{Foo: "v2p", Bar: 2}, - outputPtr: new(*unmarshalTestStruct), - expectedVal: &unmarshalTestStruct{Foo: "v2p", Bar: 2}, - }, - { - name: "unmarshal nil pointer", - inputValue: (*unmarshalTestStruct)(nil), - outputPtr: &struct{ v *unmarshalTestStruct }{v: &unmarshalTestStruct{}}, // placeholder to be replaced - expectedVal: (*unmarshalTestStruct)(nil), - }, - { - name: "convertible types", - inputValue: int32(42), - outputPtr: new(int64), - expectedVal: int64(42), - }, - { - name: "pointer to pointer destination", - inputValue: 12345, - outputPtr: new(*int), - expectedVal: ptr(12345), - }, - { - name: "unmarshal to any", - inputValue: unmarshalTestStruct{Foo: "any", Bar: 101}, - outputPtr: new(any), - expectedVal: unmarshalTestStruct{Foo: "any", Bar: 101}, - }, - } +func TestSerialization_PointerTypes(t *testing.T) { + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + t.Run("pointer to struct", func(t *testing.T) { + input := &myStruct{A: "test"} + data, err := s.Marshal(input) + require.NoError(t, err) - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - data, err := s.Marshal(tc.inputValue) + var result *myStruct + err = s.Unmarshal(data, &result) require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "test", result.A) + }) - // Special handling for the nil test case to correctly pass the pointer. - if tc.name == "unmarshal nil pointer" { - target := tc.outputPtr.(*struct{ v *unmarshalTestStruct }) - err = s.Unmarshal(data, &target.v) - require.NoError(t, err) - assert.Nil(t, target.v) - return + t.Run("double pointer", func(t *testing.T) { + value := &myStruct{A: "double"} + input := &value + data, err := s.Marshal(input) + require.NoError(t, err) + + var result **myStruct + err = s.Unmarshal(data, &result) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, *result) + assert.Equal(t, "double", (*result).A) + }) + + t.Run("slice of pointers", func(t *testing.T) { + input := []*myStruct{ + {A: "first"}, + {A: "second"}, } + data, err := s.Marshal(input) + require.NoError(t, err) - err = s.Unmarshal(data, tc.outputPtr) + var result []*myStruct + err = s.Unmarshal(data, &result) require.NoError(t, err) + require.Len(t, result, 2) + assert.Equal(t, "first", result[0].A) + assert.Equal(t, "second", result[1].A) + }) - // Dereference the pointer to get the actual value for comparison. - actualVal := reflect.ValueOf(tc.outputPtr).Elem().Interface() - assert.Equal(t, tc.expectedVal, actualVal) + t.Run("map with pointer to pointer values", func(t *testing.T) { + v1 := &myStruct{A: "v1"} + v2 := &myStruct{A: "v2"} + input := map[string]**myStruct{ + "a": &v1, + "b": &v2, + } + data, err := s.Marshal(input) + require.NoError(t, err) + + var result map[string]**myStruct + err = s.Unmarshal(data, &result) + require.NoError(t, err) + require.Len(t, result, 2) + assert.Equal(t, "v1", (**result["a"]).A) + assert.Equal(t, "v2", (**result["b"]).A) }) - } - }) - - t.Run("error cases", func(t *testing.T) { - data, err := s.Marshal(123) - require.NoError(t, err) - - t.Run("destination not a pointer", func(t *testing.T) { - var output int - err := s.Unmarshal(data, output) - require.Error(t, err) - assert.Contains(t, err.Error(), "value must be a non-nil pointer") }) + } +} - t.Run("destination is a nil pointer", func(t *testing.T) { - var output *int // nil - err := s.Unmarshal(data, output) - require.Error(t, err) - assert.Contains(t, err.Error(), "value must be a non-nil pointer") - }) +func TestSerialization_InterfaceTypes(t *testing.T) { + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + t.Run("interface value", func(t *testing.T) { + var input myInterface = &myStruct{A: "interface"} + data, err := s.Marshal(input) + require.NoError(t, err) - t.Run("type mismatch", func(t *testing.T) { - strData, mErr := s.Marshal("i am a string") - require.NoError(t, mErr) + var result myInterface + err = s.Unmarshal(data, &result) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, "interface", result.(*myStruct).A) + }) - var output int - err := s.Unmarshal(strData, &output) - require.Error(t, err) - assert.Contains(t, err.Error(), "cannot assign") - }) + t.Run("slice of interfaces with nil", func(t *testing.T) { + input := []myInterface{ + nil, + &myStruct{A: "first"}, + &myStruct{A: "second"}, + } + data, err := s.Marshal(input) + require.NoError(t, err) - t.Run("unconvertible types", func(t *testing.T) { - intData, mErr := s.Marshal(123) - require.NoError(t, mErr) + var result []myInterface + err = s.Unmarshal(data, &result) + require.NoError(t, err) + require.Len(t, result, 3) + assert.Nil(t, result[0]) + assert.Equal(t, "first", result[1].(*myStruct).A) + assert.Equal(t, "second", result[2].(*myStruct).A) + }) - var output bool - err := s.Unmarshal(intData, &output) - require.Error(t, err) - assert.Contains(t, err.Error(), "cannot assign") + t.Run("map with interface values and nil", func(t *testing.T) { + input := map[string]myInterface{ + "nil": nil, + "value": &myStruct{A: "test"}, + } + data, err := s.Marshal(input) + require.NoError(t, err) + + var result map[string]myInterface + err = s.Unmarshal(data, &result) + require.NoError(t, err) + require.Len(t, result, 2) + assert.Nil(t, result["nil"]) + assert.Equal(t, "test", result["value"].(*myStruct).A) + }) }) - }) + } +} + +func TestSerialization_MapWithStructKeys(t *testing.T) { + for serializerName, s := range getSerializers() { + t.Run(serializerName, func(t *testing.T) { + input := map[myStruct]int{ + {A: "key1"}: 100, + {A: "key2"}: 200, + } + data, err := s.Marshal(input) + require.NoError(t, err) + + var result map[myStruct]int + err = s.Unmarshal(data, &result) + require.NoError(t, err) + assert.Equal(t, 100, result[myStruct{A: "key1"}]) + assert.Equal(t, 200, result[myStruct{A: "key2"}]) + }) + } } diff --git a/schema/agentic_message.go b/schema/agentic_message.go index 2474e5df8..fe9647f38 100644 --- a/schema/agentic_message.go +++ b/schema/agentic_message.go @@ -986,6 +986,11 @@ func ConcatAgenticMessages(msgs []*AgenticMessage) (*AgenticMessage, error) { for _, idx := range blockIndices { blocks = append(blocks, indexToBlock[idx]) } + } else if len(blocks) > 1 { + blocks, err = concatAdjacentFunctionToolResultBlocks(blocks) + if err != nil { + return nil, err + } } if len(extraList) > 0 { @@ -1642,12 +1647,134 @@ func concatFunctionToolResults(results []*FunctionToolResult) (*FunctionToolResu return nil, fmt.Errorf("expected tool name '%s' for function tool result, but got '%s'", ret.Name, r.Name) } - for _, b := range r.Content { - if b == nil { - continue + var err error + ret.Content, err = concatFunctionToolResultContent(ret.Content, r.Content) + if err != nil { + return nil, err + } + } + + return ret, nil +} + +func concatAdjacentFunctionToolResultBlocks(blocks []*ContentBlock) ([]*ContentBlock, error) { + if len(blocks) <= 1 { + return blocks, nil + } + + ret := make([]*ContentBlock, 0, len(blocks)) + for _, block := range blocks { + if len(ret) == 0 || !canConcatFunctionToolResultBlocks(ret[len(ret)-1], block) { + ret = append(ret, block) + continue + } + + merged, err := concatFunctionToolResultBlocks(ret[len(ret)-1], block) + if err != nil { + return nil, err + } + ret[len(ret)-1] = merged + } + + return ret, nil +} + +func canConcatFunctionToolResultBlocks(a, b *ContentBlock) bool { + if a == nil || b == nil || + a.Type != ContentBlockTypeFunctionToolResult || + b.Type != ContentBlockTypeFunctionToolResult || + a.FunctionToolResult == nil || + b.FunctionToolResult == nil { + return false + } + + if a.FunctionToolResult.CallID != "" && b.FunctionToolResult.CallID != "" && + a.FunctionToolResult.CallID != b.FunctionToolResult.CallID { + return false + } + if a.FunctionToolResult.Name != "" && b.FunctionToolResult.Name != "" && + a.FunctionToolResult.Name != b.FunctionToolResult.Name { + return false + } + + return a.FunctionToolResult.CallID != "" || b.FunctionToolResult.CallID != "" || + (a.FunctionToolResult.Name != "" && a.FunctionToolResult.Name == b.FunctionToolResult.Name) +} + +func concatFunctionToolResultBlocks(a, b *ContentBlock) (*ContentBlock, error) { + result, err := concatFunctionToolResults([]*FunctionToolResult{a.FunctionToolResult, b.FunctionToolResult}) + if err != nil { + return nil, err + } + + block := NewContentBlock(result) + var extras []map[string]any + if len(a.Extra) > 0 { + extras = append(extras, a.Extra) + } + if len(b.Extra) > 0 { + extras = append(extras, b.Extra) + } + if len(extras) > 0 { + block.Extra, err = concatExtra(extras) + if err != nil { + return nil, fmt.Errorf("failed to concat function tool result block extras: %w", err) + } + } + + return block, nil +} + +func concatFunctionToolResultContent( + left, right []*FunctionToolResultContentBlock, +) ([]*FunctionToolResultContentBlock, error) { + ret := append([]*FunctionToolResultContentBlock(nil), left...) + for _, block := range right { + if block == nil { + continue + } + if len(ret) > 0 && canConcatFunctionToolResultTextBlocks(ret[len(ret)-1], block) { + merged, err := concatFunctionToolResultTextBlocks(ret[len(ret)-1], block) + if err != nil { + return nil, err } - ret.Content = append(ret.Content, b) + ret[len(ret)-1] = merged + continue + } + ret = append(ret, block) + } + + return ret, nil +} + +func canConcatFunctionToolResultTextBlocks(a, b *FunctionToolResultContentBlock) bool { + return a != nil && b != nil && + a.Type == FunctionToolResultContentBlockTypeText && + b.Type == FunctionToolResultContentBlockTypeText && + a.Text != nil && b.Text != nil +} + +func concatFunctionToolResultTextBlocks( + a, b *FunctionToolResultContentBlock, +) (*FunctionToolResultContentBlock, error) { + ret := &FunctionToolResultContentBlock{ + Type: FunctionToolResultContentBlockTypeText, + Text: &UserInputText{Text: a.Text.Text + b.Text.Text}, + } + + var extras []map[string]any + if len(a.Extra) > 0 { + extras = append(extras, a.Extra) + } + if len(b.Extra) > 0 { + extras = append(extras, b.Extra) + } + if len(extras) > 0 { + extra, err := concatExtra(extras) + if err != nil { + return nil, fmt.Errorf("failed to concat function tool result content extras: %w", err) } + ret.Extra = extra } return ret, nil diff --git a/schema/agentic_message_test.go b/schema/agentic_message_test.go index 32dc96c2e..c5980bf00 100644 --- a/schema/agentic_message_test.go +++ b/schema/agentic_message_test.go @@ -526,9 +526,52 @@ func TestConcatAgenticMessages(t *testing.T) { assert.Len(t, result.ContentBlocks, 1) assert.Equal(t, "call_123", result.ContentBlocks[0].FunctionToolResult.CallID) assert.Equal(t, "get_weather", result.ContentBlocks[0].FunctionToolResult.Name) - assert.Equal(t, 2, len(result.ContentBlocks[0].FunctionToolResult.Content)) - assert.Equal(t, `{"temp`, result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) - assert.Equal(t, `":72}`, result.ContentBlocks[0].FunctionToolResult.Content[1].Text.Text) + assert.Equal(t, 1, len(result.ContentBlocks[0].FunctionToolResult.Content)) + assert.Equal(t, `{"temp":72}`, result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) + }) + + t.Run("concat function tool result without streaming meta", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeFunctionToolResult, + FunctionToolResult: &FunctionToolResult{ + CallID: "call_stream", + Name: "execute", + Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: "first\n"}}, + }, + }, + }, + }, + }, + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeFunctionToolResult, + FunctionToolResult: &FunctionToolResult{ + CallID: "call_stream", + Name: "execute", + Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: "second\n"}}, + }, + }, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + require.NotNil(t, result.ContentBlocks[0].FunctionToolResult) + assert.Equal(t, "call_stream", result.ContentBlocks[0].FunctionToolResult.CallID) + assert.Equal(t, "execute", result.ContentBlocks[0].FunctionToolResult.Name) + require.Len(t, result.ContentBlocks[0].FunctionToolResult.Content, 1) + assert.Equal(t, "first\nsecond\n", result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) }) t.Run("concat server tool call", func(t *testing.T) { @@ -1726,4 +1769,19 @@ func TestConcatFunctionToolResults(t *testing.T) { assert.Equal(t, "hello", got.Content[0].Text.Text) assert.Equal(t, "http://img.png", got.Content[1].Image.URL) }) + + t.Run("text chunks", func(t *testing.T) { + results := []*FunctionToolResult{ + {CallID: "c1", Name: "tool1", Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: "hello "}}, + }}, + {CallID: "c1", Name: "tool1", Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: "world"}}, + }}, + } + got, err := concatFunctionToolResults(results) + require.NoError(t, err) + require.Len(t, got.Content, 1) + assert.Equal(t, "hello world", got.Content[0].Text.Text) + }) } diff --git a/schema/serialization.go b/schema/serialization.go index 169bf9ee9..d379ddb4b 100644 --- a/schema/serialization.go +++ b/schema/serialization.go @@ -29,6 +29,8 @@ func init() { RegisterName[[]*Message]("_eino_message_slice") RegisterName[*AgenticMessage]("_eino_agentic_message") RegisterName[[]*AgenticMessage]("_eino_agentic_message_slice") + RegisterName[*ToolInfo]("_eino_tool_info") + RegisterName[[]*ToolInfo]("_eino_tool_info_slice") RegisterName[Document]("_eino_document") RegisterName[RoleType]("_eino_role_type") RegisterName[ToolCall]("_eino_tool_call") @@ -146,3 +148,34 @@ func Register[T any]() { panic(err) } } + +// Serializer encodes and decodes persisted Eino values. +type Serializer interface { + Marshal(v any) ([]byte, error) + Unmarshal(data []byte, v any) error +} + +// HumanReadableSerializer produces clean, human-readable JSON output for serialization. +// It can be used with compose.WithSerializer() to store checkpoints in a human-readable format. +// +// Unlike the default InternalSerializer which uses verbose wrapper structures for type preservation, +// HumanReadableSerializer produces clean JSON that: +// - Uses standard JSON field names from struct tags +// - Omits empty fields when `omitempty` is specified +// - Only adds "$type" annotations for custom registered types stored in interface{} fields +// - Produces significantly smaller output for most use cases +// +// Example usage: +// +// graph, err := compose.NewGraph[Input, Output]( +// compose.WithCheckPointStore(store), +// compose.WithSerializer(&schema.HumanReadableSerializer{}), +// ) +// +// Note: All custom types stored in interface{} fields must be registered using +// schema.RegisterName[T]() or schema.Register[T]() for proper deserialization. +type HumanReadableSerializer = serialization.HumanReadableSerializer + +// GobSerializer serializes values using Go's encoding/gob package. +// It can be used with compose.WithSerializer and other serializer hooks. +type GobSerializer = serialization.GobSerializer diff --git a/schema/serialization_test.go b/schema/serialization_test.go index d17cc4092..dc601ec67 100644 --- a/schema/serialization_test.go +++ b/schema/serialization_test.go @@ -153,7 +153,6 @@ func TestRegister(t *testing.T) { }() Register[[]int]() - Register[map[string]any]() Register[[]*testStruct1]() Register[[]testStruct1]()