From b48412f8fa814ed81b22dea23a6b2056a7ead2fa Mon Sep 17 00:00:00 2001 From: buthim Date: Tue, 7 Jul 2026 12:28:48 +0800 Subject: [PATCH 1/6] feat: add visual A2A delegation host tools Co-Authored-By: Claude Opus 4.8 --- docs/a2a.md | 91 ++- docs/api-reference.md | 17 + docs/streaming.md | 46 +- examples/streaming-chat-copilotkit/README.md | 40 ++ .../skillruntime/profile_reconcile_test.go | 20 +- mcp_types.go | 138 +++++ pkg/bridges/sse/handler.go | 18 +- pkg/bridges/sse/handler_test.go | 97 +++ pkg/bridges/subagentstream/mux.go | 212 +++++++ pkg/bridges/subagentstream/mux_test.go | 241 ++++++++ pkg/hosttools/a2adelegation/delegator.go | 408 ++++++++++++ pkg/hosttools/a2adelegation/delegator_test.go | 583 ++++++++++++++++++ pkg/hosttools/a2adelegation/events.go | 105 ++++ .../a2adelegation/import_boundary_test.go | 61 ++ pkg/hosttools/a2adelegation/mapping.go | 309 ++++++++++ pkg/hosttools/a2adelegation/mcpserver.go | 273 ++++++++ pkg/hosttools/a2adelegation/types.go | 238 +++++++ runner.go | 10 +- runtime_admin_test.go | 165 ++++- 19 files changed, 3049 insertions(+), 23 deletions(-) create mode 100644 pkg/bridges/subagentstream/mux.go create mode 100644 pkg/bridges/subagentstream/mux_test.go create mode 100644 pkg/hosttools/a2adelegation/delegator.go create mode 100644 pkg/hosttools/a2adelegation/delegator_test.go create mode 100644 pkg/hosttools/a2adelegation/events.go create mode 100644 pkg/hosttools/a2adelegation/import_boundary_test.go create mode 100644 pkg/hosttools/a2adelegation/mapping.go create mode 100644 pkg/hosttools/a2adelegation/mcpserver.go create mode 100644 pkg/hosttools/a2adelegation/types.go diff --git a/docs/a2a.md b/docs/a2a.md index 99313c6..d926e72 100644 --- a/docs/a2a.md +++ b/docs/a2a.md @@ -143,6 +143,95 @@ _, err = client.CancelTask(ctx, a2a.CancelTaskRequest{ }) ``` +## Visual Delegation Host Tools + +The optional host-owned visual subagent delegation layer sits above the A2A client layer: + +- `pkg/hosttools/a2adelegation` exposes a curated registry, `delegate_to_agent` + MCP tool server, delegation event bus, A2A event mapper, and `Delegator`. +- `pkg/bridges/subagentstream` overlays delegation events onto an existing + parent AG-UI stream as `CUSTOM` events named `subagent.*`. +- `pkg/bridges/sse.Options.SubagentBus` enables the overlay in the stock SSE + handler for AG-UI callers. + +The model-facing tool input accepts registry keys only: + +```json +{ + "agent": "research", + "objective": "Find current A2A streaming behavior.", + "input": { + "prompt": "Summarize exact APIs and artifact behavior.", + "context": "We are implementing visual delegation." + }, + "constraints": { + "timeout_seconds": 180, + "stream": true, + "max_artifacts": 10 + } +} +``` + +`endpoint_url`, unknown top-level fields, and unknown nested input/constraint +fields are rejected. Host code owns the registry entry (`RemoteAgentSpec`), auth, +tenant, timeout policy, artifact limits, accepted output modes, and trusted +origins. `constraints.max_artifacts`, when supplied, limits only the final +`DelegationResult.Artifacts` returned to the parent model; live `subagent.*` +artifact events remain a UI side channel. The parent model receives only the +final structured `DelegationResult` JSON through the MCP tool result. Live remote +progress is published to the bus and rendered as AG-UI custom events: + +```json +{ + "type": "CUSTOM", + "name": "subagent.text.delta", + "value": { + "runId": "run-...", + "parentToolCallId": "tool-...", + "delegationId": "del-...", + "agentKey": "research", + "remoteProtocol": "a2a", + "remoteTaskId": "task-...", + "delta": "Searching official examples..." + } +} +``` + +`parentToolCallId` is included only when the host can supply the parent provider +tool-call ID, for example by setting `a2adelegation.MCPServerOptions.ParentToolCallID` +or by publishing `DelegationEvent` values with that field. The MCP JSON-RPC +request ID is not treated as the parent model tool-call ID. UI grouping should +tolerate this field being absent and fall back to `runId` + `delegationId`. + +Runtime-created MCP sidecars can be injected into a run without asking the model +for run IDs, URLs, or credentials. A `RuntimeServiceManager` returns a +`RuntimeServiceRef` with metadata keys: `agentadaptor.mcp.enabled=true`, +`agentadaptor.mcp.key`, `agentadaptor.mcp.transport`, `agentadaptor.mcp.url`, +`agentadaptor.mcp.command`, `agentadaptor.mcp.args_json`, +`agentadaptor.mcp.env_json`, `agentadaptor.mcp.headers_json`, +`agentadaptor.mcp.bearer_token_env_var`, `agentadaptor.mcp.required`, and +`agentadaptor.mcp.required_reason`. If `agentadaptor.mcp.key` is omitted, the SDK +falls back to the runtime ref `Name`, then `ID`. If transport is omitted, it +defaults to `stdio` when a command is present and `http` otherwise. Non-stdio +servers use `RuntimeServiceRef.URL` when `agentadaptor.mcp.url` is omitted; +stdio servers use `RuntimeServiceRef.Command` when `agentadaptor.mcp.command` is +omitted. The SDK appends those servers after runtime ensure, before profile +materialization, so adapter MCP config and session fingerprints include the +per-run delegation endpoint. + +Security boundary: concrete adapters stay unaware of A2A delegation; remote +protocol dumps are not appended to `RunResult.RawStreams`; remote artifacts are +reported as structured references/metadata unless host policy explicitly stores +and exposes content elsewhere. + ## Non-Goals -Visual subagent delegation, A2A-to-local adapter routing, bridge-managed durable task persistence, and default push delivery infrastructure are outside this slice. Issue #5 tracks the visual subagent flow above these protocol primitives. +Automatic A2A-to-local adapter routing, bridge-managed durable task persistence, +default push delivery infrastructure, dynamic remote-agent discovery, built-in +tenant/auth policy, and production per-run MCP sidecar lifecycle remain outside +this slice. + +Visual subagent delegation is an optional host-owned layer built from +`pkg/hosttools/a2adelegation`, `pkg/bridges/subagentstream`, and runtime-service +MCP injection. Core SDK execution remains protocol-agnostic and does not route to +remote A2A agents automatically. diff --git a/docs/api-reference.md b/docs/api-reference.md index 2e79cd4..47ff49e 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -66,6 +66,23 @@ Context and injection: - `WithSkills(refs...)` - `WithMCP(cfg)` - `WithProfileResources(resources)` — per-run profile desired state bundle. + +Runtime service MCP injection: a `RuntimeServiceManager` may append run-scoped MCP +servers after `RuntimeServiceManager.Ensure` by returning `RuntimeServiceRef` +values whose metadata includes `agentadaptor.mcp.enabled=true`. This does not +require an explicit `WithMCP` option. Recognized metadata keys are +`agentadaptor.mcp.enabled`, `agentadaptor.mcp.key`, +`agentadaptor.mcp.transport`, `agentadaptor.mcp.url`, +`agentadaptor.mcp.command`, `agentadaptor.mcp.args_json`, +`agentadaptor.mcp.env_json`, `agentadaptor.mcp.headers_json`, +`agentadaptor.mcp.bearer_token_env_var`, `agentadaptor.mcp.required`, and +`agentadaptor.mcp.required_reason`. The SDK converts enabled refs to +`MCPServerSpec`s, appends them to the effective MCP config, validates them with +ordinary MCP capability checks, and includes them in `DriverRunRequest.MCP` plus +the profile/session fingerprint. Malformed JSON metadata and duplicate MCP keys +fail the run before adapter start. Other runtime metadata remains runtime-service +metadata; it is not promoted into MCP server configuration. + - `WithAgents(specs...)` - `WithHooks(specs...)` - `WithProfileConfig(patches...)` diff --git a/docs/streaming.md b/docs/streaming.md index 7e2273a..d449543 100644 --- a/docs/streaming.md +++ b/docs/streaming.md @@ -145,11 +145,43 @@ data: {"type":"RUN_FINISHED",...} - `CORSAllowedOrigin`: 浏览器前端必须；空字符串不加 CORS 头 - `WriteTimeout`: 单帧写超时，默认 30s - `RunOptions`: 应用到每次请求的 RunOption 列表（例如 `WithRunPolicy`） -- `DecodeRequest`: 自定义请求体解码 +- `SubagentBus`: AG-UI 协议下把 host-published visual subagent delegation events 叠加为 `CUSTOM` 事件；Raw 协议不使用该 overlay -## 4. 配置 +## 4. Visual subagent delegation overlay -### 4.1 per-binding 默认开 +Hosts that expose a `delegate_to_agent` MCP tool can keep remote A2A progress out +of the parent model context while still rendering it live in the same AG-UI/SSE +stream. The delegation tool publishes `pkg/hosttools/a2adelegation.DelegationEvent` +values to a run-scoped bus; `pkg/bridges/subagentstream` merges that bus with the +parent run's AG-UI stream and forwards each remote event as an AG-UI `CUSTOM` +event such as `subagent.started`, `subagent.text.delta`, `subagent.artifact`, or +`subagent.finished`. + +```go +bus := a2adelegation.NewEventBus(256) + +mux.Handle("/v1/chat", sse.Handler(sdk, sse.Options{ + Protocol: sse.AGUI, + SubagentBus: bus, +})) + +// The MCP tool server for the same run calls delegator.Delegate(...), which +// returns only the final structured DelegationResult to the parent model while +// publishing live progress onto bus.SubscribeRun(ctx, handle.RunID()). +``` + +`SubagentBus` is honored by the stock SSE handler only when `Protocol: sse.AGUI`. +For Raw SSE, hosts that need the same side channel should call +`subagentstream.Wrap(...)` themselves and choose their own wire shape. + +For lower-level hosts, use `subagentstream.WrapAGUI(ctx, handle, opts)` directly +instead of `agui.WrapWithContext`. Parent AG-UI lifecycle events are preserved; +remote subagent events remain a UI side channel and are not appended to +`RunResult.Output`, `RunResult.RawStreams`, or ordinary tool output. + +## 5. 配置 + +### 5.1 per-binding 默认开 ```go sdk := agentadaptor.New( @@ -169,7 +201,7 @@ handle2, _ := sdk.Start(ctx, "batch", 覆盖顺序：per-call `WithStreaming` / `WithoutStreaming` > per-binding `WithDefaultStreaming` > 默认 off。 -### 4.2 背压 +### 5.2 背压 ```go sdk := agentadaptor.New( @@ -183,7 +215,7 @@ sdk := agentadaptor.New( **宿主务必及时 drain `StreamEvents()`**，否则 Block 模式可能挂起，Drop 模式会丢事件。 -## 5. 能力与降级 +## 6. 能力与降级 每个 streaming-aware adapter 通过 `StreamAwareDriver.StreamCapability()` 声明自己能提供什么： @@ -199,7 +231,7 @@ sdk := agentadaptor.New( - `Reasoning=false` → 不发 `REASONING_*` - HITL 事件默认映射成 AG-UI tool-call lifecycle，便于 CopilotKit 这类客户端渲染审批卡片；需要旧行为时使用 `agui.WithDecisionMode(agui.DecisionAsCustom)` -## 6. 常见问题 +## 7. 常见问题 **Q: 开了 `WithStreaming` 但 `StreamEvents()` 一直没有 payload？** A: 检查 adapter 是否实现了 `StreamAwareDriver`。未实现的 adapter 会走普通路径，`StreamEvents()` 返回 closed channel。 @@ -216,7 +248,7 @@ A: 安全。每个 `Start()` 创建独立 handle / sink / codex app-server 子 **Q: SSE 断连后 run 还在跑吗？** A: 不在。SSE handler 在 client 断连时会调 `handle.Cancel()`，ctx cancellation 穿透到 codex 子进程。 -## 7. AG-UI 前后端版本对齐 +## 8. AG-UI 前后端版本对齐 Go 侧通过 `go.mod` 固定 `github.com/ag-ui-protocol/ag-ui/sdks/community/go`；`examples/streaming-chat-copilotkit/web` 通过 `package-lock.json` 固定 `@ag-ui/core`。两边不是同一个坐标系，升级任一侧时都可能出现 Zod/Go `Validate` 行为漂移。 diff --git a/examples/streaming-chat-copilotkit/README.md b/examples/streaming-chat-copilotkit/README.md index 53fbb0f..5a9da43 100644 --- a/examples/streaming-chat-copilotkit/README.md +++ b/examples/streaming-chat-copilotkit/README.md @@ -87,6 +87,46 @@ Browser | `AGENT_BACKEND_URL` | `http://localhost:8080/agent` | CopilotRuntime 转发的 AG-UI 端点 | | `NEXT_PUBLIC_AGENT_BACKEND_BASE` | `http://localhost:8080` | 浏览器旁路请求 base URL | +## Visual subagent delegation hook + +The backend can render remote A2A delegation progress in the same CopilotKit +AG-UI stream by wiring `sse.Options.SubagentBus` (or the lower-level +`subagentstream.WrapAGUI`) to the run session. A host-owned `delegate_to_agent` +MCP server publishes `a2adelegation.DelegationEvent` values while the tool call +is blocked on the remote task. CopilotKit receives those updates as AG-UI +`CUSTOM` events named `subagent.started`, `subagent.text.delta`, +`subagent.artifact`, `subagent.finished`, etc.; Claude/Codex/Cursor receives only +the final structured MCP result. + +This example's current main path still starts one local parent CLI. To prove the +visual delegation product path, run a Codex A2A bridge separately (see +`examples/a2a-local` and `docs/a2a.md`), register its Agent Card in an +`a2adelegation.Registry`, start an HTTP `a2adelegation.MCPServer` for each parent +run, and expose that server through runtime-backed MCP metadata: + +```text +agentadaptor.mcp.enabled=true +agentadaptor.mcp.key=a2a-delegation +agentadaptor.mcp.transport=http +agentadaptor.mcp.url=http://127.0.0.1:/mcp +agentadaptor.mcp.bearer_token_env_var=A2A_DELEGATION_TOKEN +``` + +`agentadaptor.mcp.enabled=true` is the promotion switch. Other optional keys +include `agentadaptor.mcp.headers_json`, `agentadaptor.mcp.args_json`, +`agentadaptor.mcp.env_json`, `agentadaptor.mcp.required`, and +`agentadaptor.mcp.required_reason`. The SDK validates these as ordinary MCP +config; malformed JSON or duplicate MCP keys fail the run before the parent +adapter starts. + +The important boundary for UI authors: render the `subagent.*` custom events as +a nested visual group, but do not add them to the chat transcript as parent model +messages. `subagent.*` event values may include `parentToolCallId` when the host +can correlate the delegation sidecar with a parent provider tool-use ID, but UI +code should not require it. Use `delegationId` as the stable nested group key, +with `runId` as the run scope. The parent transcript should contain the normal +tool call plus the concise final delegation result only. + ## 相关文档 - [`docs/workstream-user-message-event.md`](../../docs/workstream-user-message-event.md) diff --git a/internal/skillruntime/profile_reconcile_test.go b/internal/skillruntime/profile_reconcile_test.go index cdd1240..46bb5f4 100644 --- a/internal/skillruntime/profile_reconcile_test.go +++ b/internal/skillruntime/profile_reconcile_test.go @@ -107,7 +107,7 @@ func TestReconcileProfileSkillsPreservesExternalConflict(t *testing.T) { skillsHome := filepath.Join(profileDir, "skills") managedSource := createProfileSkillDir(t, t.TempDir(), "managed-main", "managed") externalSource := createProfileSkillDir(t, t.TempDir(), "external-main", "external") - if err := os.MkdirAll(skillsHome, 0o755); err != nil { + if err := os.MkdirAll(skillsHome, 0755); err != nil { t.Fatalf("mkdir skills home: %v", err) } if err := os.Symlink(externalSource, filepath.Join(skillsHome, "main")); err != nil { @@ -129,7 +129,11 @@ func TestReconcileProfileSkillsPreservesExternalConflict(t *testing.T) { if err != nil { t.Fatalf("resolve external target: %v", err) } - if filepath.Clean(resolved) != filepath.Clean(externalSource) { + expected, err := filepath.EvalSymlinks(externalSource) + if err != nil { + t.Fatalf("resolve expected external target: %v", err) + } + if filepath.Clean(resolved) != filepath.Clean(expected) { t.Fatalf("expected external target preserved, got %s", resolved) } manifest, err := profilestate.LoadManifest(profileDir) @@ -146,7 +150,7 @@ func TestReconcileProfileSkillsRejectsExternalConflictWhenRequested(t *testing.T skillsHome := filepath.Join(profileDir, "skills") managedSource := createProfileSkillDir(t, t.TempDir(), "managed-review", "managed") externalSource := createProfileSkillDir(t, t.TempDir(), "external-review", "external") - if err := os.MkdirAll(skillsHome, 0o755); err != nil { + if err := os.MkdirAll(skillsHome, 0755); err != nil { t.Fatalf("mkdir skills home: %v", err) } if err := os.Symlink(externalSource, filepath.Join(skillsHome, "review")); err != nil { @@ -169,7 +173,11 @@ func TestReconcileProfileSkillsRejectsExternalConflictWhenRequested(t *testing.T if err != nil { t.Fatalf("resolve external target: %v", err) } - if filepath.Clean(resolved) != filepath.Clean(externalSource) { + expected, err := filepath.EvalSymlinks(externalSource) + if err != nil { + t.Fatalf("resolve expected external target: %v", err) + } + if filepath.Clean(resolved) != filepath.Clean(expected) { t.Fatalf("expected external target preserved, got %s", resolved) } } @@ -177,11 +185,11 @@ func TestReconcileProfileSkillsRejectsExternalConflictWhenRequested(t *testing.T func createProfileSkillDir(t *testing.T, root, name, body string) string { t.Helper() skillDir := filepath.Join(root, name) - if err := os.MkdirAll(skillDir, 0o755); err != nil { + if err := os.MkdirAll(skillDir, 0755); err != nil { t.Fatalf("mkdir skill dir: %v", err) } content := "---\nname: " + name + "\n---\n" + body + "\n" - if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644); err != nil { + if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0644); err != nil { t.Fatalf("write skill: %v", err) } return skillDir diff --git a/mcp_types.go b/mcp_types.go index 5f90a6e..5a49031 100644 --- a/mcp_types.go +++ b/mcp_types.go @@ -1,8 +1,10 @@ package agentadaptor import ( + "encoding/json" "fmt" "slices" + "strconv" "strings" ) @@ -108,6 +110,142 @@ func resolveMCPPayload(defaults, override *MCPConfig, caps MCPCapability) (MCPPa return prepareMCPPayload(*effective, caps) } +func resolveMCPPayloadWithRuntime(defaults, override *MCPConfig, refs []RuntimeServiceRef, caps MCPCapability) (MCPPayload, error) { + effective := MCPConfig{} + if defaults != nil { + effective.Servers = cloneMCPServerSpecs(defaults.Servers) + } + if override != nil { + effective.Servers = cloneMCPServerSpecs(override.Servers) + } + + runtimeServers, err := mcpServersFromRuntimeRefs(refs) + if err != nil { + return MCPPayload{}, err + } + effective.Servers = append(effective.Servers, runtimeServers...) + if len(effective.Servers) == 0 { + return MCPPayload{}, nil + } + return prepareMCPPayload(effective, caps) +} + +func mcpServersFromRuntimeRefs(refs []RuntimeServiceRef) ([]MCPServerSpec, error) { + if len(refs) == 0 { + return nil, nil + } + out := make([]MCPServerSpec, 0, len(refs)) + for _, ref := range refs { + spec, ok, err := mcpServerFromRuntimeRef(ref) + if err != nil { + return nil, err + } + if ok { + out = append(out, spec) + } + } + return out, nil +} + +func mcpServerFromRuntimeRef(ref RuntimeServiceRef) (MCPServerSpec, bool, error) { + metadata := ref.Metadata + if !runtimeMCPEnabled(metadata) { + return MCPServerSpec{}, false, nil + } + + spec := MCPServerSpec{ + Key: strings.TrimSpace(metadata["agentadaptor.mcp.key"]), + Transport: MCPTransport(strings.TrimSpace(metadata["agentadaptor.mcp.transport"])), + Command: strings.TrimSpace(metadata["agentadaptor.mcp.command"]), + URL: strings.TrimSpace(metadata["agentadaptor.mcp.url"]), + BearerTokenEnvVar: strings.TrimSpace(metadata["agentadaptor.mcp.bearer_token_env_var"]), + RequiredReason: strings.TrimSpace(metadata["agentadaptor.mcp.required_reason"]), + } + if spec.Key == "" { + spec.Key = strings.TrimSpace(ref.Name) + } + if spec.Key == "" { + spec.Key = strings.TrimSpace(ref.ID) + } + if spec.Transport == "" { + if spec.Command != "" { + spec.Transport = MCPTransportStdio + } else { + spec.Transport = MCPTransportHTTP + } + } + if spec.URL == "" && spec.Transport != MCPTransportStdio { + spec.URL = strings.TrimSpace(ref.URL) + } + if spec.Command == "" && spec.Transport == MCPTransportStdio { + spec.Command = strings.TrimSpace(ref.Command) + } + + var err error + if spec.Args, err = parseRuntimeMCPStringSlice(metadata, "agentadaptor.mcp.args_json"); err != nil { + return MCPServerSpec{}, false, err + } + if spec.Env, err = parseRuntimeMCPStringMap(metadata, "agentadaptor.mcp.env_json"); err != nil { + return MCPServerSpec{}, false, err + } + if spec.Headers, err = parseRuntimeMCPStringMap(metadata, "agentadaptor.mcp.headers_json"); err != nil { + return MCPServerSpec{}, false, err + } + if spec.Required, err = parseRuntimeMCPBool(metadata, "agentadaptor.mcp.required"); err != nil { + return MCPServerSpec{}, false, err + } + + return spec, true, nil +} + +func runtimeMCPEnabled(metadata map[string]string) bool { + if len(metadata) == 0 { + return false + } + raw, ok := metadata["agentadaptor.mcp.enabled"] + if !ok { + return false + } + enabled, err := strconv.ParseBool(strings.TrimSpace(raw)) + return err == nil && enabled +} + +func parseRuntimeMCPStringSlice(metadata map[string]string, key string) ([]string, error) { + raw := strings.TrimSpace(metadata[key]) + if raw == "" { + return nil, nil + } + var values []string + if err := json.Unmarshal([]byte(raw), &values); err != nil { + return nil, fmt.Errorf("%w: runtime MCP metadata %s must be a JSON string array: %v", ErrInvalidMCPConfig, key, err) + } + return cloneStrings(values), nil +} + +func parseRuntimeMCPStringMap(metadata map[string]string, key string) (map[string]string, error) { + raw := strings.TrimSpace(metadata[key]) + if raw == "" { + return nil, nil + } + var values map[string]string + if err := json.Unmarshal([]byte(raw), &values); err != nil { + return nil, fmt.Errorf("%w: runtime MCP metadata %s must be a JSON object with string values: %v", ErrInvalidMCPConfig, key, err) + } + return cloneStringMap(values), nil +} + +func parseRuntimeMCPBool(metadata map[string]string, key string) (bool, error) { + raw := strings.TrimSpace(metadata[key]) + if raw == "" { + return false, nil + } + value, err := strconv.ParseBool(raw) + if err != nil { + return false, fmt.Errorf("%w: runtime MCP metadata %s must be boolean: %v", ErrInvalidMCPConfig, key, err) + } + return value, nil +} + func prepareMCPPayload(cfg MCPConfig, caps MCPCapability) (MCPPayload, error) { if len(cfg.Servers) == 0 { return MCPPayload{Fingerprint: stableHash("mcp", []MCPServerSpec{})}, nil diff --git a/pkg/bridges/sse/handler.go b/pkg/bridges/sse/handler.go index 4e74f24..b4d6a13 100644 --- a/pkg/bridges/sse/handler.go +++ b/pkg/bridges/sse/handler.go @@ -44,6 +44,7 @@ import ( agentadaptor "github.com/agent-dance/agent-adaptor" "github.com/agent-dance/agent-adaptor/pkg/bridges/agui" + "github.com/agent-dance/agent-adaptor/pkg/bridges/subagentstream" ) // Protocol selects the full on-wire contract (inbound + outbound). @@ -88,6 +89,12 @@ type Options struct { // They apply before any options derived from the request body, so // the body can still override them (e.g. session binding). RunOptions []agentadaptor.RunOption + + // SubagentBus, when set, overlays host-published visual subagent + // delegation events onto the AG-UI response stream as CUSTOM events. + // Parent adapter events are left unchanged; the remote stream remains a + // host-side UI side channel and is not fed into the parent model context. + SubagentBus subagentstream.EventBus } // RawRequest is the canonical Raw-protocol chat request body. @@ -161,7 +168,7 @@ func Handler(sdk agentadaptor.SDK, opts Options) http.Handler { } ctx := r.Context() - err = streamEvents(ctx, writer, w, handle, opts.Protocol) + err = streamEvents(ctx, writer, w, handle, opts) if err != nil && !errors.Is(err, context.Canceled) { _ = writer.WriteErrorEvent(ctx, w, err, "") } @@ -235,10 +242,13 @@ func decodeRawRequest(r *http.Request) (*RawRequest, error) { // streamEvents drains the handle's streams into SSE frames according to // the chosen protocol. It returns when the stream is exhausted or the // context is cancelled. -func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, w io.Writer, handle agentadaptor.RunHandle, protocol Protocol) error { - switch protocol { +func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, w io.Writer, handle agentadaptor.RunHandle, opts Options) error { + switch opts.Protocol { case AGUI: out := agui.WrapWithContext(ctx, handle) + if opts.SubagentBus != nil { + out = subagentstream.WrapAGUI(ctx, handle, subagentstream.MuxOptions{Bus: opts.SubagentBus}) + } for ev := range out { select { case <-ctx.Done(): @@ -253,7 +263,7 @@ func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, w io.Writer, h case Raw: return streamRaw(ctx, w, handle) default: - return fmt.Errorf("sse: unknown protocol %d", protocol) + return fmt.Errorf("sse: unknown protocol %d", opts.Protocol) } } diff --git a/pkg/bridges/sse/handler_test.go b/pkg/bridges/sse/handler_test.go index 10f6ba8..7b45f45 100644 --- a/pkg/bridges/sse/handler_test.go +++ b/pkg/bridges/sse/handler_test.go @@ -13,6 +13,7 @@ import ( agentadaptor "github.com/agent-dance/agent-adaptor" "github.com/agent-dance/agent-adaptor/memory" "github.com/agent-dance/agent-adaptor/pkg/bridges/sse" + "github.com/agent-dance/agent-adaptor/pkg/hosttools/a2adelegation" ) // fakeAdapter emits a fixed StreamPayload sequence so we can exercise the @@ -285,6 +286,102 @@ func TestSSEHandlerRejectsWrongMethod(t *testing.T) { } } +type blockingAdapter struct{} + +func (blockingAdapter) Descriptor() agentadaptor.DriverDescriptor { + return agentadaptor.DriverDescriptor{Type: "blocking", DisplayName: "Blocking"} +} +func (blockingAdapter) ValidateConfig(any) error { return nil } +func (blockingAdapter) StreamCapability() agentadaptor.StreamCapability { + return agentadaptor.StreamCapability{Native: true, TokenLevel: true} +} +func (blockingAdapter) Run(ctx context.Context, req agentadaptor.DriverRunRequest, sink agentadaptor.EventSink) (agentadaptor.DriverRunResult, error) { + if err := sink.EmitStream(agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunStarted, ThreadID: "t", RunID: req.RunID}); err != nil { + return agentadaptor.DriverRunResult{}, err + } + select { + case <-ctx.Done(): + return agentadaptor.DriverRunResult{}, ctx.Err() + case <-time.After(200 * time.Millisecond): + } + if err := sink.EmitStream(agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunFinished, ThreadID: "t", RunID: req.RunID}); err != nil { + return agentadaptor.DriverRunResult{}, err + } + return agentadaptor.DriverRunResult{Output: "done", ExitCode: 0}, nil +} + +func newBlockingSDK(t *testing.T) agentadaptor.SDK { + t.Helper() + return agentadaptor.New( + agentadaptor.WithDefaultAgent(agentadaptor.Bind(blockingAdapter{}, nil)), + agentadaptor.WithSessionStore(memory.NewSessionStore()), + ) +} + +func TestSSEHandlerAGUIOverlaysSubagentEvents(t *testing.T) { + t.Parallel() + sdk := newBlockingSDK(t) + bus := scriptedSubagentBus{} + srv := httptest.NewServer(sse.Handler(sdk, sse.Options{Protocol: sse.AGUI, SubagentBus: bus})) + defer srv.Close() + + body := strings.NewReader(`{ + "threadId": "t-1", + "runId": "r-1", + "messages": [{"id":"m-1","role":"user","content":"hi"}] + }`) + resp, err := http.Post(srv.URL, "application/json", body) + if err != nil { + t.Fatalf("post: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status: %d", resp.StatusCode) + } + + frames := readSSEFrames(t, resp.Body, 2*time.Second) + seenSubagent := false + for _, f := range frames { + if f.data == "" { + continue + } + var payload map[string]any + if err := json.Unmarshal([]byte(f.data), &payload); err != nil { + t.Fatalf("decode frame data: %v (raw=%q)", err, f.data) + } + if payload["type"] == "CUSTOM" && payload["name"] == "subagent.started" { + seenSubagent = true + value, _ := payload["value"].(map[string]any) + if value["delegationId"] != "del-1" || value["agentKey"] != "research" { + t.Fatalf("unexpected subagent payload: %#v", value) + } + } + } + if !seenSubagent { + t.Fatalf("subagent custom event not observed in frames: %#v", frames) + } +} + +type scriptedSubagentBus struct{} + +func (scriptedSubagentBus) SubscribeRun(ctx context.Context, runID string) <-chan a2adelegation.DelegationEvent { + out := make(chan a2adelegation.DelegationEvent, 2) + go func() { + defer close(out) + select { + case <-ctx.Done(): + return + case out <- a2adelegation.DelegationEvent{RunID: runID, DelegationID: "del-1", AgentKey: "research", Kind: a2adelegation.DelegationStarted}: + } + select { + case <-ctx.Done(): + return + case out <- a2adelegation.DelegationEvent{RunID: runID, DelegationID: "del-1", AgentKey: "research", Kind: a2adelegation.DelegationFinished, Status: "completed"}: + } + }() + return out +} + // --------------------------------------------------------------------------- // SSE parsing helpers // --------------------------------------------------------------------------- diff --git a/pkg/bridges/subagentstream/mux.go b/pkg/bridges/subagentstream/mux.go new file mode 100644 index 0000000..bdae351 --- /dev/null +++ b/pkg/bridges/subagentstream/mux.go @@ -0,0 +1,212 @@ +package subagentstream + +import ( + "context" + "sync/atomic" + "time" + + aguievents "github.com/ag-ui-protocol/ag-ui/sdks/community/go/pkg/core/events" + + agentadaptor "github.com/agent-dance/agent-adaptor" + "github.com/agent-dance/agent-adaptor/pkg/bridges/agui" + "github.com/agent-dance/agent-adaptor/pkg/hosttools/a2adelegation" +) + +type EventBus interface { + SubscribeRun(ctx context.Context, runID string) <-chan a2adelegation.DelegationEvent +} + +type MuxOptions struct { + Bus EventBus +} + +type Event struct { + ID uint64 + AGUI aguievents.Event + Raw *agentadaptor.StreamPayload + Subagent *a2adelegation.DelegationEvent +} + +func WrapAGUI(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) <-chan aguievents.Event { + out := make(chan aguievents.Event, 32) + go func() { + defer close(out) + for ev := range Wrap(ctx, handle, opts) { + if ev.AGUI == nil { + continue + } + select { + case <-ctx.Done(): + return + case out <- ev.AGUI: + } + } + }() + return out +} + +func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) <-chan Event { + out := make(chan Event, 64) + var seq atomic.Uint64 + + sendAGUI := func(ev aguievents.Event) bool { + wrapped := Event{ID: seq.Add(1), AGUI: ev} + select { + case <-ctx.Done(): + return false + case out <- wrapped: + return true + } + } + sendSubagent := func(ev a2adelegation.DelegationEvent) bool { + aguiEvent := AGUICustomEvent(ev) + evCopy := ev + wrapped := Event{ID: seq.Add(1), AGUI: aguiEvent, Subagent: &evCopy} + select { + case <-ctx.Done(): + return false + case out <- wrapped: + return true + } + } + drainSubagents := func(subagents <-chan a2adelegation.DelegationEvent) bool { + for subagents != nil { + select { + case ev, ok := <-subagents: + if !ok { + return true + } + if !sendSubagent(ev) { + return false + } + default: + return true + } + } + return true + } + + go func() { + defer close(out) + parent := agui.WrapWithContext(ctx, handle) + var subagents <-chan a2adelegation.DelegationEvent + if opts.Bus != nil { + subagents = opts.Bus.SubscribeRun(ctx, handle.RunID()) + } + for parent != nil || subagents != nil { + select { + case <-ctx.Done(): + cancelHandle(handle) + return + case ev, ok := <-parent: + if !ok { + parent = nil + if !drainSubagents(subagents) { + return + } + subagents = nil + continue + } + if isTerminalAGUI(ev) { + if !drainSubagents(subagents) { + return + } + if !sendAGUI(ev) { + return + } + parent = nil + subagents = nil + continue + } + if !sendAGUI(ev) { + return + } + case ev, ok := <-subagents: + if !ok { + subagents = nil + continue + } + if !sendSubagent(ev) { + return + } + } + } + }() + return out +} + +func cancelHandle(handle agentadaptor.RunHandle) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = handle.Cancel(ctx) +} + +func isTerminalAGUI(ev aguievents.Event) bool { + if ev == nil { + return false + } + switch ev.Type() { + case aguievents.EventTypeRunFinished, aguievents.EventTypeRunError: + return true + default: + return false + } +} + +func AGUICustomEvent(ev a2adelegation.DelegationEvent) aguievents.Event { + value := map[string]any{ + "runId": ev.RunID, + "parentToolCallId": ev.ParentToolCallID, + "delegationId": ev.DelegationID, + "agentKey": ev.AgentKey, + "agentName": ev.AgentName, + "remoteProtocol": ev.Protocol, + "remoteTaskId": ev.RemoteTaskID, + "remoteContextId": ev.RemoteContextID, + "messageId": ev.RemoteMessageID, + "artifactId": ev.RemoteArtifactID, + "delta": ev.Delta, + "text": ev.Text, + "status": ev.Status, + } + if ev.Artifact != nil { + value["artifact"] = ev.Artifact + } + if ev.Error != nil { + value["error"] = ev.Error + } + if ev.Raw != nil { + value["raw"] = ev.Raw + } + for key, val := range value { + if val == "" || val == nil { + delete(value, key) + } + } + return aguievents.NewCustomEvent(string(ev.Kind), aguievents.WithValue(value)) +} + +func StreamPayload(ev a2adelegation.DelegationEvent) agentadaptor.StreamPayload { + return agentadaptor.StreamPayload{ + Kind: "", + Name: string(ev.Kind), + RunID: ev.RunID, + ToolCallID: ev.ParentToolCallID, + MessageID: ev.RemoteMessageID, + Delta: ev.Delta, + Raw: map[string]any{ + "delegation_id": ev.DelegationID, + "agent_key": ev.AgentKey, + "agent_name": ev.AgentName, + "parent_tool_call_id": ev.ParentToolCallID, + "remote_protocol": ev.Protocol, + "remote_task_id": ev.RemoteTaskID, + "remote_context_id": ev.RemoteContextID, + "remote_message_id": ev.RemoteMessageID, + "remote_artifact_id": ev.RemoteArtifactID, + "delta": ev.Delta, + "text": ev.Text, + "status": ev.Status, + }, + } +} diff --git a/pkg/bridges/subagentstream/mux_test.go b/pkg/bridges/subagentstream/mux_test.go new file mode 100644 index 0000000..9a2eafc --- /dev/null +++ b/pkg/bridges/subagentstream/mux_test.go @@ -0,0 +1,241 @@ +package subagentstream_test + +import ( + "context" + "sync" + "testing" + "time" + + aguievents "github.com/ag-ui-protocol/ag-ui/sdks/community/go/pkg/core/events" + + agentadaptor "github.com/agent-dance/agent-adaptor" + "github.com/agent-dance/agent-adaptor/pkg/bridges/subagentstream" + "github.com/agent-dance/agent-adaptor/pkg/hosttools/a2adelegation" +) + +func TestAGUICustomEventMapsDelegationFields(t *testing.T) { + t.Parallel() + ev := subagentstream.AGUICustomEvent(a2adelegation.DelegationEvent{ + RunID: "run-1", + ParentToolCallID: "tool-1", + DelegationID: "del-1", + AgentKey: "research", + AgentName: "Research", + Protocol: a2adelegation.ProtocolA2A, + RemoteTaskID: "task-1", + RemoteContextID: "ctx-1", + RemoteMessageID: "msg-1", + Kind: a2adelegation.DelegationTextDelta, + Delta: "hello", + }) + custom, ok := ev.(*aguievents.CustomEvent) + if !ok { + t.Fatalf("expected CustomEvent, got %T", ev) + } + if custom.Name != string(a2adelegation.DelegationTextDelta) { + t.Fatalf("custom name: got %q", custom.Name) + } + value := custom.Value.(map[string]any) + if value["delegationId"] != "del-1" || value["delta"] != "hello" || value["remoteTaskId"] != "task-1" { + t.Fatalf("unexpected custom value: %#v", value) + } + if _, ok := value["raw"]; ok { + t.Fatalf("raw should be omitted when empty: %#v", value) + } +} + +func TestStreamPayloadUsesCustomPassThroughShape(t *testing.T) { + t.Parallel() + payload := subagentstream.StreamPayload(a2adelegation.DelegationEvent{ + RunID: "run-1", + DelegationID: "del-1", + AgentKey: "research", + Kind: a2adelegation.DelegationStatus, + Status: "working", + }) + if payload.Kind != "" || payload.Name != string(a2adelegation.DelegationStatus) { + t.Fatalf("expected custom StreamPayload, got %#v", payload) + } + if payload.Raw["delegation_id"] != "del-1" || payload.Raw["status"] != "working" { + t.Fatalf("unexpected raw payload: %#v", payload.Raw) + } +} + +func TestWrapMergesParentAndDelegationAGUIEvents(t *testing.T) { + t.Parallel() + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload, 4), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := a2adelegation.NewEventBus(8) + out := subagentstream.Wrap(context.Background(), handle, subagentstream.MuxOptions{Bus: bus}) + + handle.stream <- agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunStarted, ThreadID: "thread-1", RunID: "run-1"} + events := collectMuxEvents(t, out, 1) + bus.Publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "del-1", AgentKey: "research", Kind: a2adelegation.DelegationStarted}) + close(handle.stream) + close(handle.done) + events = append(events, collectMuxEvents(t, out, 1)...) + + if events[0].ID == 0 || events[1].ID <= events[0].ID { + t.Fatalf("expected monotonic mux IDs, got %#v", events) + } + if events[0].AGUI.Type() != aguievents.EventTypeRunStarted { + t.Fatalf("first event should be parent RUN_STARTED, got %s", events[0].AGUI.Type()) + } + if events[1].AGUI.Type() != aguievents.EventTypeCustom || events[1].Subagent == nil { + t.Fatalf("second event should be subagent custom, got %#v", events[1]) + } +} + +func TestWrapDrainsBufferedSubagentsBeforeParentTerminal(t *testing.T) { + t.Parallel() + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload, 1), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := a2adelegation.NewEventBus(8) + bus.Publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "del-1", AgentKey: "research", Kind: a2adelegation.DelegationFinished, Status: "completed"}) + out := subagentstream.Wrap(context.Background(), handle, subagentstream.MuxOptions{Bus: bus}) + + handle.stream <- agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunFinished, ThreadID: "thread-1", RunID: "run-1"} + close(handle.stream) + close(handle.done) + events := collectMuxEvents(t, out, 3) + seenSubagent := false + for _, ev := range events[:len(events)-1] { + if ev.AGUI.Type() == aguievents.EventTypeCustom && ev.Subagent != nil { + seenSubagent = true + } + } + if !seenSubagent { + t.Fatalf("expected buffered subagent before terminal, got %#v", events) + } + if events[len(events)-1].AGUI.Type() != aguievents.EventTypeRunFinished { + t.Fatalf("terminal should remain last, got %#v", events) + } +} + +func TestWrapAGUIExitsWhenContextCanceledAndConsumerStops(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload, 128), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + out := subagentstream.WrapAGUI(ctx, handle, subagentstream.MuxOptions{}) + for i := 0; i < 128; i++ { + handle.stream <- agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunStarted, ThreadID: "thread-1", RunID: "run-1"} + } + cancel() + close(handle.stream) + close(handle.done) + for { + select { + case _, ok := <-out: + if !ok { + return + } + case <-time.After(time.Second): + t.Fatal("WrapAGUI did not close after cancellation") + } + } +} + +func TestWrapCancelUsesBestEffortContext(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + out := subagentstream.Wrap(ctx, handle, subagentstream.MuxOptions{}) + cancel() + for range out { + } + if !handle.sawLiveCancelContext() { + t.Fatalf("expected at least one Cancel call with a live context, got errs=%v", handle.cancelErrs()) + } +} + +func collectMuxEvents(t *testing.T, ch <-chan subagentstream.Event, want int) []subagentstream.Event { + t.Helper() + out := make([]subagentstream.Event, 0, want) + for len(out) < want { + select { + case ev, ok := <-ch: + if !ok { + t.Fatalf("mux closed after %d events, wanted %d", len(out), want) + } + out = append(out, ev) + case <-time.After(time.Second): + t.Fatalf("timeout waiting for mux event %d/%d", len(out), want) + } + } + return out +} + +type fakeHandle struct { + stream chan agentadaptor.StreamPayload + events chan agentadaptor.RunEvent + done chan struct{} + runID string + runResult agentadaptor.RunResult + runErr error + + mu sync.Mutex + cancelCtxErrs []error +} + +func (f *fakeHandle) Events() <-chan agentadaptor.RunEvent { return f.events } +func (f *fakeHandle) StreamEvents() <-chan agentadaptor.StreamPayload { return f.stream } +func (f *fakeHandle) RunID() string { return f.runID } +func (f *fakeHandle) Cancel(ctx context.Context) error { + f.mu.Lock() + defer f.mu.Unlock() + f.cancelCtxErrs = append(f.cancelCtxErrs, ctx.Err()) + return nil +} +func (f *fakeHandle) DecisionRequests() <-chan agentadaptor.DecisionRequest { + ch := make(chan agentadaptor.DecisionRequest) + close(ch) + return ch +} +func (f *fakeHandle) ResolveDecision(string, agentadaptor.DecisionResponse) error { + return agentadaptor.ErrRunEnded +} +func (f *fakeHandle) Wait(ctx context.Context) (agentadaptor.RunResult, error) { + select { + case <-ctx.Done(): + return agentadaptor.RunResult{}, ctx.Err() + case <-f.done: + return f.runResult, f.runErr + } +} + +func (f *fakeHandle) cancelErrs() []error { + f.mu.Lock() + defer f.mu.Unlock() + return append([]error(nil), f.cancelCtxErrs...) +} + +func (f *fakeHandle) sawLiveCancelContext() bool { + for _, err := range f.cancelErrs() { + if err == nil { + return true + } + } + return false +} diff --git a/pkg/hosttools/a2adelegation/delegator.go b/pkg/hosttools/a2adelegation/delegator.go new file mode 100644 index 0000000..3775c31 --- /dev/null +++ b/pkg/hosttools/a2adelegation/delegator.go @@ -0,0 +1,408 @@ +package a2adelegation + +import ( + "context" + "encoding/json" + "fmt" + "io" + "strings" + "time" + + clienta2a "github.com/agent-dance/agent-adaptor/pkg/clients/a2a" +) + +type A2AClient interface { + AgentCard(ctx context.Context) (clienta2a.AgentCard, error) + Send(ctx context.Context, req clienta2a.SendRequest) (clienta2a.Task, error) + SendStream(ctx context.Context, req clienta2a.SendRequest) (*clienta2a.Stream, error) + GetTask(ctx context.Context, req clienta2a.GetTaskRequest) (clienta2a.Task, error) + CancelTask(ctx context.Context, req clienta2a.CancelTaskRequest) (clienta2a.Task, error) +} + +type ClientFactory func(RemoteAgentSpec) A2AClient + +type Delegator struct { + Registry *Registry + Bus *EventBus + NewClient ClientFactory + NewID func() string +} + +func NewDelegator(registry *Registry, bus *EventBus) *Delegator { + return &Delegator{Registry: registry, Bus: bus} +} + +func (d *Delegator) Delegate(ctx context.Context, req DelegationRequest) (DelegationResult, error) { + if d == nil || d.Registry == nil { + return DelegationResult{}, &DelegationError{Code: "configuration_error", Message: "delegation registry is required"} + } + spec, ok := d.Registry.Lookup(req.Agent) + if !ok { + err := &DelegationError{Code: "agent_not_found", Message: fmt.Sprintf("remote agent %q is not registered", req.Agent)} + return DelegationResult{DelegationID: d.newID(), Agent: req.Agent, RemoteProtocol: ProtocolA2A, Status: "failed", Error: err}, err + } + delegationID := d.newID() + if delegationID == "" { + delegationID = "del-unknown" + } + baseEvent := DelegationEvent{ + RunID: req.RunID, + ParentToolCallID: req.ParentToolCallID, + DelegationID: delegationID, + AgentKey: spec.Key, + AgentName: displayName(spec), + Protocol: ProtocolA2A, + } + baseResult := DelegationResult{DelegationID: delegationID, Agent: spec.Key, RemoteProtocol: ProtocolA2A, Status: "running"} + + client := d.clientFor(spec) + card, err := client.AgentCard(ctx) + if err != nil && spec.AgentCard == nil { + derr := &DelegationError{Code: "agent_unavailable", Message: err.Error(), Retryable: true} + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + if spec.AgentCard != nil { + card = *spec.AgentCard + } + if spec.Policy.RequireStreaming && !card.Capabilities.Streaming { + derr := &DelegationError{Code: "capability_unsupported", Message: "remote agent does not advertise streaming"} + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + + timeout := clampTimeout(req.Timeout, spec.Policy.MaxTimeout) + if timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + + message, err := messageForDelegation(req) + if err != nil { + derr := delegationErr(err) + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + send := clienta2a.SendRequest{ + Message: message, + Tenant: spec.Tenant, + AcceptedOutputModes: spec.AcceptedOutputModes, + Metadata: cloneAnyMap(req.Metadata), + } + if req.Tenant != "" { + send.Tenant = req.Tenant + } + if req.Stream || card.Capabilities.Streaming { + return d.delegateStreaming(ctx, client, spec, send, baseEvent, baseResult, req.MaxArtifacts) + } + return d.delegatePolling(ctx, client, spec, send, baseEvent, baseResult, req.MaxArtifacts) +} + +func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spec RemoteAgentSpec, send clienta2a.SendRequest, baseEvent DelegationEvent, baseResult DelegationResult, maxArtifacts *int) (DelegationResult, error) { + stream, err := client.SendStream(ctx, send) + if err != nil { + if spec.Policy.RequireStreaming { + derr := &DelegationError{Code: "stream_unavailable", Message: err.Error(), Retryable: true} + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + return d.delegatePolling(ctx, client, spec, send, baseEvent, baseResult, maxArtifacts) + } + defer stream.Close() + mapper := newEventMapper(baseEvent) + var lastTask clienta2a.Task + for { + event, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + if lastTask.ID != "" { + if recovered, ok := d.recoverTask(ctx, client, spec, lastTask.ID); ok { + lastTask = recovered + for _, ev := range mapper.taskEvents(recovered) { + d.publish(ev) + } + return d.finishTask(baseEvent, baseResult, recovered, spec.Policy, maxArtifacts) + } + } + derr := &DelegationError{Code: "stream_interrupted", Message: err.Error(), Retryable: true} + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + if event.Task != nil { + lastTask = *event.Task + } + for _, ev := range mapper.Map(event) { + d.publish(ev) + } + if event.Task != nil && executionFinalState(event.Task.Status.State) { + return d.finishTask(baseEvent, baseResult, *event.Task, spec.Policy, maxArtifacts) + } + if event.Status != nil && executionFinalState(event.Status.State) { + task, ok := d.recoverTask(ctx, client, spec, event.TaskID) + if ok { + return d.finishTask(baseEvent, baseResult, task, spec.Policy, maxArtifacts) + } + baseResult.RemoteTaskID = event.TaskID + baseResult.RemoteContextID = event.ContextID + baseResult.Status = statusFromState(event.Status.State) + if derr := policyErrorForState(event.Status.State, spec.Policy); derr != nil { + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + if baseResult.Status != "completed" { + baseResult.Error = &DelegationError{Code: errorCodeFromState(event.Status.State), Message: "remote task did not complete successfully", RemoteStatus: string(event.Status.State)} + } + return baseResult, terminalError(event.Status.State, spec.Policy) + } + } + if lastTask.ID != "" && executionFinalState(lastTask.Status.State) { + return d.finishTask(baseEvent, baseResult, lastTask, spec.Policy, maxArtifacts) + } + derr := &DelegationError{Code: "stream_interrupted", Message: "remote stream ended before terminal state", Retryable: true} + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr +} + +func (d *Delegator) delegatePolling(ctx context.Context, client A2AClient, spec RemoteAgentSpec, send clienta2a.SendRequest, baseEvent DelegationEvent, baseResult DelegationResult, maxArtifacts *int) (DelegationResult, error) { + send.ReturnImmediately = true + task, err := client.Send(ctx, send) + if err != nil { + derr := &DelegationError{Code: "agent_unavailable", Message: err.Error(), Retryable: true} + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + mapper := newEventMapper(baseEvent) + d.publish(mapper.Started(task.ID, task.ContextID)) + for _, ev := range mapper.taskEvents(task) { + d.publish(ev) + } + if executionFinalState(task.Status.State) { + return d.finishTask(baseEvent, baseResult, task, spec.Policy, maxArtifacts) + } + interval := spec.Policy.PollInterval + if interval <= 0 { + interval = 500 * time.Millisecond + } + maxPolls := spec.Policy.MaxPolls + if maxPolls <= 0 { + maxPolls = 600 + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + for i := 0; i < maxPolls; i++ { + select { + case <-ctx.Done(): + d.cancelRemote(context.Background(), client, task.ID, send.Tenant, baseEvent) + derr := &DelegationError{Code: "cancelled", Message: ctx.Err().Error(), Retryable: true} + baseResult.Status = "cancelled" + baseResult.Error = derr + return baseResult, derr + case <-ticker.C: + } + task, err = client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: task.ID, Tenant: send.Tenant}) + if err != nil { + continue + } + for _, ev := range mapper.taskEvents(task) { + d.publish(ev) + } + if executionFinalState(task.Status.State) { + return d.finishTask(baseEvent, baseResult, task, spec.Policy, maxArtifacts) + } + } + d.cancelRemote(context.Background(), client, task.ID, send.Tenant, baseEvent) + derr := &DelegationError{Code: "remote_timeout", Message: "remote task did not finish before timeout", Retryable: true, RemoteStatus: string(task.Status.State)} + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr +} + +func (d *Delegator) recoverTask(ctx context.Context, client A2AClient, spec RemoteAgentSpec, taskID string) (clienta2a.Task, bool) { + if taskID == "" { + return clienta2a.Task{}, false + } + task, err := client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: taskID, Tenant: spec.Tenant}) + return task, err == nil && executionFinalState(task.Status.State) +} + +func (d *Delegator) cancelRemote(ctx context.Context, client A2AClient, taskID, tenant string, base DelegationEvent) { + if taskID != "" { + _, _ = client.CancelTask(ctx, clienta2a.CancelTaskRequest{TaskID: taskID, Tenant: tenant, Metadata: map[string]any{"reason": "parent_cancelled", "delegation_id": base.DelegationID}}) + } + ev := base + ev.Kind = DelegationCancelled + ev.RemoteTaskID = taskID + ev.Status = "cancelled" + d.publish(ev) +} + +func (d *Delegator) publish(ev DelegationEvent) { + if d.Bus != nil { + d.Bus.Publish(ev) + } +} + +func (d *Delegator) clientFor(spec RemoteAgentSpec) A2AClient { + if d.NewClient != nil { + return d.NewClient(spec) + } + return clienta2a.New(clienta2a.Options{ + AgentCardURL: spec.AgentCardURL, + Auth: spec.Auth, + HTTPClient: spec.HTTPClient, + TrustedAuthOrigins: spec.TrustedAuthOrigins, + AcceptedOutputModes: spec.AcceptedOutputModes, + PreferredTransports: spec.PreferredTransports, + }) +} + +func (d *Delegator) newID() string { + if d != nil && d.NewID != nil { + return d.NewID() + } + return "del-" + time.Now().UTC().Format("20060102150405.000000000") +} + +func messageForDelegation(req DelegationRequest) (clienta2a.Message, error) { + parts := []clienta2a.Part{{Kind: clienta2a.PartText, Text: promptFor(req)}} + for _, artifact := range req.Artifacts { + uri := strings.TrimSpace(artifact.URI) + if uri == "" { + return clienta2a.Message{}, &DelegationError{Code: "invalid_artifact", Message: "input artifact uri is required"} + } + parts = append(parts, clienta2a.Part{Kind: clienta2a.PartURL, URL: uri, MediaType: artifact.MediaType, Filename: artifact.Name}) + } + return clienta2a.Message{Role: "user", Parts: parts}, nil +} + +func promptFor(req DelegationRequest) string { + parts := []string{strings.TrimSpace(req.Objective), strings.TrimSpace(req.Prompt), strings.TrimSpace(req.Context)} + out := []string{} + for _, part := range parts { + if part != "" { + out = append(out, part) + } + } + return strings.Join(out, "\n\n") +} + +func displayName(spec RemoteAgentSpec) string { + if spec.DisplayName != "" { + return spec.DisplayName + } + if spec.AgentCard != nil && spec.AgentCard.Name != "" { + return spec.AgentCard.Name + } + return spec.Key +} + +func clampTimeout(requested, max time.Duration) time.Duration { + if requested <= 0 { + return max + } + if max > 0 && requested > max { + return max + } + return requested +} + +func failedEvent(base DelegationEvent, derr *DelegationError) DelegationEvent { + ev := base + ev.Kind = DelegationFailed + if derr != nil && (derr.Code == "cancelled" || derr.Code == "remote_cancelled") { + ev.Kind = DelegationCancelled + } + ev.Error = derr + if derr != nil { + ev.Status = derr.RemoteStatus + } + return ev +} + +func (d *Delegator) finishTask(baseEvent DelegationEvent, baseResult DelegationResult, task clienta2a.Task, policy DelegationPolicy, maxArtifacts *int) (DelegationResult, error) { + if derr := policyErrorForTask(task, policy); derr != nil { + d.publish(failedEvent(baseEvent, derr)) + baseResult.RemoteTaskID = task.ID + baseResult.RemoteContextID = task.ContextID + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } + result := resultFromTask(baseResult, task) + if task.Status.State == clienta2a.TaskStateInputRequired && policy.AllowInputRequired { + result.Error = nil + } + result.Artifacts = limitArtifacts(result.Artifacts, maxArtifacts) + return result, terminalError(task.Status.State, policy) +} + +func policyErrorForTask(task clienta2a.Task, policy DelegationPolicy) *DelegationError { + if derr := policyErrorForState(task.Status.State, policy); derr != nil { + return derr + } + if policy.MaxArtifactBytes <= 0 { + return nil + } + for _, artifact := range task.Artifacts { + if artifactKnownBytes(artifact) > policy.MaxArtifactBytes { + return &DelegationError{Code: "artifact_too_large", Message: "remote artifact exceeds max artifact byte policy", RemoteStatus: string(task.Status.State), Metadata: map[string]any{"artifact_id": artifact.ID, "max_bytes": policy.MaxArtifactBytes}} + } + } + return nil +} + +func policyErrorForState(state clienta2a.TaskState, policy DelegationPolicy) *DelegationError { + if state == clienta2a.TaskStateInputRequired && !policy.AllowInputRequired { + return &DelegationError{Code: "input_required", Message: "remote task requires input", RemoteStatus: string(state)} + } + return nil +} + +func artifactKnownBytes(artifact clienta2a.Artifact) int64 { + var total int64 + for _, part := range artifact.Parts { + total += int64(len(part.Text) + len(part.Raw) + len(part.URL) + len(part.MediaType) + len(part.Filename)) + if part.Data != nil { + if raw, err := json.Marshal(part.Data); err == nil { + total += int64(len(raw)) + } + } + } + return total +} + +func limitArtifacts(artifacts []DelegationArtifact, max *int) []DelegationArtifact { + if max == nil || *max >= len(artifacts) { + return artifacts + } + if *max <= 0 { + return nil + } + return artifacts[:*max] +} + +func terminalError(state clienta2a.TaskState, policy DelegationPolicy) error { + if state == clienta2a.TaskStateCompleted || (state == clienta2a.TaskStateInputRequired && policy.AllowInputRequired) { + return nil + } + return &DelegationError{Code: errorCodeFromState(state), Message: "remote task did not complete successfully", RemoteStatus: string(state)} +} diff --git a/pkg/hosttools/a2adelegation/delegator_test.go b/pkg/hosttools/a2adelegation/delegator_test.go new file mode 100644 index 0000000..472fb34 --- /dev/null +++ b/pkg/hosttools/a2adelegation/delegator_test.go @@ -0,0 +1,583 @@ +package a2adelegation + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + bridgea2a "github.com/agent-dance/agent-adaptor/pkg/bridges/a2a" + clienta2a "github.com/agent-dance/agent-adaptor/pkg/clients/a2a" +) + +func TestRegistryRejectsDuplicateKeys(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + _, err := NewRegistry( + RemoteAgentSpec{Key: "research", AgentCard: &card}, + RemoteAgentSpec{Key: "research", AgentCard: &card}, + ) + if err == nil { + t.Fatal("expected duplicate registry key to fail") + } +} + +func TestParseToolInputRejectsArbitraryEndpointURL(t *testing.T) { + t.Parallel() + _, err := ParseToolInput([]byte(`{"agent":"research","objective":"do work","endpoint_url":"https://evil.example/a2a"}`)) + if err == nil { + t.Fatal("expected endpoint_url to be rejected") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "invalid_tool_input" { + t.Fatalf("expected invalid_tool_input, got %T %[1]v", err) + } +} + +func TestParseToolInputRejectsUnknownNestedFields(t *testing.T) { + t.Parallel() + cases := []string{ + `{"agent":"research","objective":"do work","input":{"prompt":"hi","extra":true}}`, + `{"agent":"research","objective":"do work","input":{"artifacts":[{"uri":"file://a","extra":true}]}}`, + `{"agent":"research","objective":"do work","constraints":{"stream":true,"extra":true}}`, + } + for _, raw := range cases { + if _, err := ParseToolInput([]byte(raw)); err == nil { + t.Fatalf("expected unknown nested field to fail for %s", raw) + } + } +} + +func TestParseToolInputTracksExplicitMaxArtifacts(t *testing.T) { + t.Parallel() + input, err := ParseToolInput([]byte(`{"agent":"research","objective":"do work","constraints":{"max_artifacts":0}}`)) + if err != nil { + t.Fatalf("parse explicit max_artifacts: %v", err) + } + if !input.Constraints.MaxArtifactsSet || input.Constraints.MaxArtifacts != 0 { + t.Fatalf("expected explicit max_artifacts=0, got %#v", input.Constraints) + } + input, err = ParseToolInput([]byte(`{"agent":"research","objective":"do work"}`)) + if err != nil { + t.Fatalf("parse omitted max_artifacts: %v", err) + } + if input.Constraints.MaxArtifactsSet { + t.Fatalf("max_artifacts should be omitted: %#v", input.Constraints) + } +} + +func TestToolSchemaAllowsOnlyRegistryKeyObjectiveInputAndConstraints(t *testing.T) { + t.Parallel() + schema := ToolSchema() + props, ok := schema["properties"].(map[string]any) + if !ok { + t.Fatalf("schema properties missing: %#v", schema) + } + if _, ok := props["endpoint_url"]; ok { + t.Fatal("schema must not expose arbitrary endpoint_url") + } + if schema["additionalProperties"] != false { + t.Fatalf("expected additionalProperties=false, got %#v", schema["additionalProperties"]) + } +} + +func TestEventMapperMapsA2ABridgeArtifactsAndTerminal(t *testing.T) { + t.Parallel() + mapper := newEventMapper(DelegationEvent{RunID: "run-1", DelegationID: "del-1", AgentKey: "research"}) + events := mapper.Map(clienta2a.Event{ + Kind: clienta2a.EventArtifact, + TaskID: "task-1", + ContextID: "ctx-1", + Artifact: &clienta2a.Artifact{ + ID: "assistant-output", + Name: bridgea2a.ArtifactAssistantOutput, + Parts: []clienta2a.Part{{Kind: clienta2a.PartText, Text: "hello"}}, + }, + }) + if len(events) != 2 { + t.Fatalf("expected synthesized start + text delta, got %#v", events) + } + if events[0].Kind != DelegationStarted || events[1].Kind != DelegationTextDelta || events[1].Delta != "hello" { + t.Fatalf("unexpected mapped events: %#v", events) + } + + terminal := mapper.Map(clienta2a.Event{ + Kind: clienta2a.EventTerminal, + TaskID: "task-1", + ContextID: "ctx-1", + Status: &clienta2a.TaskStatus{State: clienta2a.TaskStateInputRequired}, + }) + if len(terminal) != 1 || terminal[0].Kind != DelegationInputRequired { + t.Fatalf("input-required terminal mapping: %#v", terminal) + } +} + +func TestEventMapperPreservesArtifactURI(t *testing.T) { + t.Parallel() + mapper := newEventMapper(DelegationEvent{RunID: "run-1", DelegationID: "del-1", AgentKey: "research"}) + events := mapper.Map(clienta2a.Event{ + Kind: clienta2a.EventArtifact, + TaskID: "task-1", + ContextID: "ctx-1", + Artifact: &clienta2a.Artifact{ + ID: "notes", + Name: "notes.md", + Parts: []clienta2a.Part{{Kind: clienta2a.PartURL, URL: "file:///notes.md", MediaType: "text/markdown"}}, + }, + }) + if len(events) != 2 { + t.Fatalf("expected started + artifact events, got %#v", events) + } + artifact := events[1].Artifact + if artifact == nil || artifact.URI != "file:///notes.md" || artifact.MediaType != "text/markdown" { + t.Fatalf("expected artifact URI/media type to be preserved, got %#v", artifact) + } +} + +func TestResultFromTaskUsesAgentAdaptorResultArtifact(t *testing.T) { + t.Parallel() + result := resultFromTask(DelegationResult{DelegationID: "del-1", Agent: "research", RemoteProtocol: ProtocolA2A}, clienta2a.Task{ + ID: "task-1", + ContextID: "ctx-1", + Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}, + Artifacts: []clienta2a.Artifact{ + { + ID: bridgea2a.ArtifactAgentAdaptorResult, + Name: bridgea2a.ArtifactAgentAdaptorResult, + Parts: []clienta2a.Part{{Kind: clienta2a.PartData, Data: map[string]any{ + "summary": "structured summary", + "result": map[string]any{"confidence": "high"}, + }}}, + }, + {ID: "notes", Name: "notes.md", Description: "notes", Parts: []clienta2a.Part{{Kind: clienta2a.PartURL, URL: "file:///notes.md", MediaType: "text/markdown"}}}, + }, + }) + if result.Status != "completed" || result.Summary != "structured summary" { + t.Fatalf("unexpected result summary/status: %#v", result) + } + if result.Metadata["confidence"] != "high" { + t.Fatalf("expected structured result metadata, got %#v", result.Metadata) + } + if len(result.Artifacts) != 1 || result.Artifacts[0].Name != "notes.md" { + t.Fatalf("expected non-internal artifact reference, got %#v", result.Artifacts) + } + if result.Artifacts[0].URI != "file:///notes.md" { + t.Fatalf("expected artifact URI to be preserved, got %#v", result.Artifacts[0]) + } +} + +func TestDelegatorPollingHappyPathEmitsEventsAndStructuredResult(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Tenant: "tenant-a", Policy: DelegationPolicy{PollInterval: time.Millisecond, MaxPolls: 3}}) + if err != nil { + t.Fatalf("registry: %v", err) + } + bus := NewEventBus(16) + client := &fakeA2AClient{ + card: card, + sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateWorking}}, + getTasks: []clienta2a.Task{{ + ID: "task-1", + ContextID: "ctx-1", + Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}, + Messages: []clienta2a.Message{{Role: "agent", Parts: []clienta2a.Part{{Kind: clienta2a.PartText, Text: "done"}}}}, + Artifacts: []clienta2a.Artifact{{ID: "notes", Name: "notes.md", Parts: []clienta2a.Part{{Kind: clienta2a.PartText, MediaType: "text/markdown"}}}}, + }}, + } + delegator := NewDelegator(registry, bus) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this"}) + if err != nil { + t.Fatalf("delegate: %v", err) + } + if result.Status != "completed" || result.RemoteTaskID != "task-1" || result.Summary != "done" { + t.Fatalf("unexpected delegation result: %#v", result) + } + if client.lastSend.Tenant != "tenant-a" || client.lastSend.Message.Parts[0].Text != "research this" { + t.Fatalf("unexpected send request: %#v", client.lastSend) + } + replayed := drainBus(t, bus, "run-1", 5) + if replayed[0].Kind != DelegationStarted || replayed[len(replayed)-1].Kind != DelegationFinished { + t.Fatalf("unexpected bus replay: %#v", replayed) + } +} + +func TestDelegatorRequireStreamingRejectsSendStreamFailure(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Policy: DelegationPolicy{RequireStreaming: true, PollInterval: time.Millisecond, MaxPolls: 1}}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{card: card, sendTask: clienta2a.Task{ID: "task-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}}} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Stream: true}) + if err == nil { + t.Fatal("expected streaming failure") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "stream_unavailable" { + t.Fatalf("expected stream_unavailable error, got %T %[1]v", err) + } + if result.Status != "failed" || client.streamCalls != 1 || client.sendCalls != 0 { + t.Fatalf("expected failed result without polling fallback, result=%#v streamCalls=%d sendCalls=%d", result, client.streamCalls, client.sendCalls) + } +} + +func TestDelegatorInputRequiredPolicy(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + for _, tc := range []struct { + name string + allow bool + wantErr bool + wantStatus string + }{ + {name: "default rejects", allow: false, wantErr: true, wantStatus: "failed"}, + {name: "policy allows", allow: true, wantErr: false, wantStatus: "input_required"}, + } { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Policy: DelegationPolicy{AllowInputRequired: tc.allow}}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{card: card, sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateInputRequired}}} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this"}) + if tc.wantErr && err == nil { + t.Fatal("expected input-required error") + } + if !tc.wantErr && err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.Status != tc.wantStatus { + t.Fatalf("status: got %q want %q result=%#v", result.Status, tc.wantStatus, result) + } + if !tc.wantErr && result.Error != nil { + t.Fatalf("allowed input-required should not set error: %#v", result.Error) + } + }) + } +} + +func TestDelegatorMaxArtifactBytesPolicy(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Policy: DelegationPolicy{MaxArtifactBytes: 4}}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{card: card, sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}, Artifacts: []clienta2a.Artifact{{ID: "large", Name: "large.txt", Parts: []clienta2a.Part{{Kind: clienta2a.PartText, Text: "too large"}}}}}} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this"}) + if err == nil { + t.Fatal("expected artifact size policy error") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "artifact_too_large" { + t.Fatalf("expected artifact_too_large error, got %T %[1]v", err) + } + if result.Status != "failed" { + t.Fatalf("expected failed result, got %#v", result) + } +} + +func TestDelegatorSendsInputArtifactsAsURLParts(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{card: card, sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}}} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + _, err = delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Artifacts: []InputArtifact{{Name: "notes.md", URI: "file:///notes.md", MediaType: "text/markdown"}}}) + if err != nil { + t.Fatalf("delegate: %v", err) + } + parts := client.lastSend.Message.Parts + if len(parts) != 2 || parts[1].Kind != clienta2a.PartURL || parts[1].URL != "file:///notes.md" || parts[1].Filename != "notes.md" || parts[1].MediaType != "text/markdown" { + t.Fatalf("expected URL artifact part, got %#v", parts) + } +} + +func TestDelegatorRejectsInputArtifactWithoutURI(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return &fakeA2AClient{card: card} } + delegator.NewID = func() string { return "del-1" } + + _, err = delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Artifacts: []InputArtifact{{Name: "notes.md"}}}) + if err == nil { + t.Fatal("expected invalid artifact error") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "invalid_artifact" { + t.Fatalf("expected invalid_artifact error, got %T %[1]v", err) + } +} + +func TestDelegatorMaxArtifactsLimitsFinalResultOnly(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{card: card, sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}, Artifacts: []clienta2a.Artifact{{ID: "a", Name: "a.txt"}, {ID: "b", Name: "b.txt"}}}} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + max := 1 + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", MaxArtifacts: &max}) + if err != nil { + t.Fatalf("delegate: %v", err) + } + if len(result.Artifacts) != 1 || result.Artifacts[0].ID != "a" { + t.Fatalf("expected final artifacts to be limited to first artifact, got %#v", result.Artifacts) + } +} + +func TestDelegatorPollingCancellationCascadesToRemoteTask(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research"} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Tenant: "tenant-a", Policy: DelegationPolicy{PollInterval: time.Millisecond, MaxPolls: 10}}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{ + card: card, + sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateWorking}}, + getErr: context.DeadlineExceeded, + } + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err = delegator.Delegate(ctx, DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this"}) + if err == nil { + t.Fatal("expected cancellation error") + } + if client.cancelCalls != 1 || client.lastCancel.TaskID != "task-1" || client.lastCancel.Tenant != "tenant-a" { + t.Fatalf("expected remote cancel cascade, calls=%d req=%#v", client.cancelCalls, client.lastCancel) + } +} + +func TestEventBusDeduplicatesTerminalEventsAndReplays(t *testing.T) { + t.Parallel() + bus := NewEventBus(8) + if !bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationStarted}) { + t.Fatal("expected started event to publish") + } + if !bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationFinished}) { + t.Fatal("expected first terminal event to publish") + } + if bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationFailed}) { + t.Fatal("duplicate terminal event should be suppressed") + } + replayed := drainBus(t, bus, "run-1", 2) + if len(replayed) != 2 || replayed[1].Kind != DelegationFinished { + t.Fatalf("unexpected replayed events: %#v", replayed) + } +} + +func TestEventBusReplayAboveSubscriberBufferDoesNotDeadlock(t *testing.T) { + t.Parallel() + bus := NewEventBus(64) + for i := 0; i < 40; i++ { + if !bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationStatus, Status: "working"}) { + t.Fatalf("publish %d failed", i) + } + } + replayed := drainBus(t, bus, "run-1", 40) + if len(replayed) != 40 { + t.Fatalf("expected 40 replayed events, got %d", len(replayed)) + } +} + +func TestEventBusPublishDoesNotBlockOnFullSubscriber(t *testing.T) { + t.Parallel() + bus := NewEventBus(0) + ctx, cancel := context.WithCancel(context.Background()) + ch := bus.SubscribeRun(ctx, "run-1") + for i := 0; i < subscriberBuffer; i++ { + bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationStatus}) + } + done := make(chan struct{}) + go func() { + bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationStatus}) + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("publish blocked on full subscriber") + } + cancel() + for range ch { + } + if !bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationStatus}) { + t.Fatal("publish after unsubscribe should still accept event") + } +} + +func drainBus(t *testing.T, bus *EventBus, runID string, want int) []DelegationEvent { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := bus.SubscribeRun(ctx, runID) + out := make([]DelegationEvent, 0, want) + for len(out) < want { + select { + case ev := <-ch: + out = append(out, ev) + case <-time.After(time.Second): + t.Fatalf("timeout waiting for replay event %d/%d", len(out), want) + } + } + return out +} + +func TestMCPServerDelegatesToolCallAndReturnsStructuredResult(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Tenant: "tenant-a", Policy: DelegationPolicy{PollInterval: time.Millisecond, MaxPolls: 2}}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{ + card: card, + sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}, Messages: []clienta2a.Message{{Role: "agent", Parts: []clienta2a.Part{{Kind: clienta2a.PartText, Text: "done"}}}}}, + } + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + server := httptest.NewServer(NewMCPServer(delegator, MCPServerOptions{RunID: "run-1", BearerToken: "token"}).Handler()) + defer server.Close() + rpcBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"delegate_to_agent","arguments":{"agent":"research","objective":"research this","constraints":{"stream":false}}}}` + req, err := http.NewRequest(http.MethodPost, server.URL, strings.NewReader(rpcBody)) + if err != nil { + t.Fatalf("request: %v", err) + } + req.Header.Set("Authorization", "Bearer token") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("post rpc: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status: %d", resp.StatusCode) + } + var envelope struct { + Result struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + IsError bool `json:"isError"` + } `json:"result"` + } + if err := json.NewDecoder(resp.Body).Decode(&envelope); err != nil { + t.Fatalf("decode response: %v", err) + } + if envelope.Result.IsError || len(envelope.Result.Content) != 1 { + t.Fatalf("unexpected tool response: %#v", envelope.Result) + } + var result DelegationResult + if err := json.Unmarshal([]byte(envelope.Result.Content[0].Text), &result); err != nil { + t.Fatalf("decode delegation result: %v", err) + } + if result.DelegationID != "del-1" || result.RemoteTaskID != "task-1" || result.Status != "completed" || result.Summary != "done" { + t.Fatalf("unexpected delegation result: %#v", result) + } +} + +func TestMCPServerRejectsMissingBearerToken(t *testing.T) { + t.Parallel() + server := httptest.NewServer(NewMCPServer(NewDelegator(nil, nil), MCPServerOptions{BearerToken: "token"}).Handler()) + defer server.Close() + resp, err := http.Post(server.URL, "application/json", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + if err != nil { + t.Fatalf("post rpc: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("expected unauthorized, got %d", resp.StatusCode) + } +} + +type fakeA2AClient struct { + card clienta2a.AgentCard + + sendTask clienta2a.Task + lastSend clienta2a.SendRequest + sendCalls int + + streamCalls int + + getTasks []clienta2a.Task + getErr error + + cancelCalls int + lastCancel clienta2a.CancelTaskRequest +} + +func (f *fakeA2AClient) AgentCard(context.Context) (clienta2a.AgentCard, error) { return f.card, nil } + +func (f *fakeA2AClient) Send(_ context.Context, req clienta2a.SendRequest) (clienta2a.Task, error) { + f.sendCalls++ + f.lastSend = req + return f.sendTask, nil +} + +func (f *fakeA2AClient) SendStream(context.Context, clienta2a.SendRequest) (*clienta2a.Stream, error) { + f.streamCalls++ + return nil, errors.New("stream unavailable") +} + +func (f *fakeA2AClient) GetTask(context.Context, clienta2a.GetTaskRequest) (clienta2a.Task, error) { + if f.getErr != nil { + return clienta2a.Task{}, f.getErr + } + if len(f.getTasks) == 0 { + return f.sendTask, nil + } + task := f.getTasks[0] + f.getTasks = f.getTasks[1:] + return task, nil +} + +func (f *fakeA2AClient) CancelTask(_ context.Context, req clienta2a.CancelTaskRequest) (clienta2a.Task, error) { + f.cancelCalls++ + f.lastCancel = req + return clienta2a.Task{ID: req.TaskID, Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCanceled}}, nil +} diff --git a/pkg/hosttools/a2adelegation/events.go b/pkg/hosttools/a2adelegation/events.go new file mode 100644 index 0000000..99935ed --- /dev/null +++ b/pkg/hosttools/a2adelegation/events.go @@ -0,0 +1,105 @@ +package a2adelegation + +import ( + "context" + "sync" + "time" +) + +const subscriberBuffer = 32 + +type EventBus struct { + mu sync.Mutex + subscribers map[string]map[chan DelegationEvent]struct{} + replay map[string][]DelegationEvent + replayLimit int + terminal map[string]map[string]struct{} +} + +func NewEventBus(replayLimit int) *EventBus { + if replayLimit < 0 { + replayLimit = 0 + } + return &EventBus{ + subscribers: map[string]map[chan DelegationEvent]struct{}{}, + replay: map[string][]DelegationEvent{}, + replayLimit: replayLimit, + terminal: map[string]map[string]struct{}{}, + } +} + +func (b *EventBus) Publish(ev DelegationEvent) bool { + if b == nil || ev.RunID == "" { + return false + } + if ev.Time.IsZero() { + ev.Time = time.Now().UTC() + } + b.mu.Lock() + defer b.mu.Unlock() + + if isTerminal(ev.Kind) { + if b.terminal[ev.RunID] == nil { + b.terminal[ev.RunID] = map[string]struct{}{} + } + if _, exists := b.terminal[ev.RunID][ev.DelegationID]; exists { + return false + } + b.terminal[ev.RunID][ev.DelegationID] = struct{}{} + } + + if b.replayLimit > 0 { + buf := append(b.replay[ev.RunID], ev) + if len(buf) > b.replayLimit { + buf = buf[len(buf)-b.replayLimit:] + } + b.replay[ev.RunID] = buf + } + for ch := range b.subscribers[ev.RunID] { + select { + case ch <- ev: + default: + } + } + return true +} + +func (b *EventBus) SubscribeRun(ctx context.Context, runID string) <-chan DelegationEvent { + if b == nil || runID == "" { + out := make(chan DelegationEvent) + close(out) + return out + } + b.mu.Lock() + replay := append([]DelegationEvent(nil), b.replay[runID]...) + out := make(chan DelegationEvent, len(replay)+subscriberBuffer) + for _, ev := range replay { + out <- ev + } + if b.subscribers[runID] == nil { + b.subscribers[runID] = map[chan DelegationEvent]struct{}{} + } + b.subscribers[runID][out] = struct{}{} + b.mu.Unlock() + + go func() { + <-ctx.Done() + b.mu.Lock() + delete(b.subscribers[runID], out) + if len(b.subscribers[runID]) == 0 { + delete(b.subscribers, runID) + } + b.mu.Unlock() + close(out) + }() + return out +} + +func isTerminal(kind DelegationEventKind) bool { + switch kind { + case DelegationFinished, DelegationFailed, DelegationCancelled, DelegationInputRequired: + return true + default: + return false + } +} diff --git a/pkg/hosttools/a2adelegation/import_boundary_test.go b/pkg/hosttools/a2adelegation/import_boundary_test.go new file mode 100644 index 0000000..d64e1ee --- /dev/null +++ b/pkg/hosttools/a2adelegation/import_boundary_test.go @@ -0,0 +1,61 @@ +package a2adelegation_test + +import ( + "go/parser" + "go/token" + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +func TestDelegationPackagesDoNotImportConcreteAdapters(t *testing.T) { + t.Parallel() + + repoRoot := repositoryRoot(t) + packages := []string{ + filepath.Join(repoRoot, "pkg", "hosttools", "a2adelegation"), + filepath.Join(repoRoot, "pkg", "bridges", "subagentstream"), + } + forbidden := []string{ + "github.com/agent-dance/agent-adaptor/claude", + "github.com/agent-dance/agent-adaptor/codex", + "github.com/agent-dance/agent-adaptor/cursor", + } + + fset := token.NewFileSet() + for _, root := range packages { + entries, err := os.ReadDir(root) + if err != nil { + t.Fatalf("readdir %s: %v", root, err) + } + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".go") || strings.HasSuffix(entry.Name(), "_test.go") { + continue + } + path := filepath.Join(root, entry.Name()) + file, err := parser.ParseFile(fset, path, nil, parser.ImportsOnly) + if err != nil { + t.Fatalf("parse %s: %v", path, err) + } + for _, imp := range file.Imports { + importPath := strings.Trim(imp.Path.Value, `"`) + for _, forbiddenImport := range forbidden { + if importPath == forbiddenImport || strings.HasPrefix(importPath, forbiddenImport+"/") { + t.Fatalf("%s imports forbidden concrete adapter %s", path, forbiddenImport) + } + } + } + } + } +} + +func repositoryRoot(t *testing.T) string { + t.Helper() + _, thisFile, _, ok := runtime.Caller(0) + if !ok { + t.Fatal("runtime.Caller failed") + } + return filepath.Clean(filepath.Join(filepath.Dir(thisFile), "..", "..", "..")) +} diff --git a/pkg/hosttools/a2adelegation/mapping.go b/pkg/hosttools/a2adelegation/mapping.go new file mode 100644 index 0000000..9b5fd2b --- /dev/null +++ b/pkg/hosttools/a2adelegation/mapping.go @@ -0,0 +1,309 @@ +package a2adelegation + +import ( + "encoding/json" + "strings" + + bridgea2a "github.com/agent-dance/agent-adaptor/pkg/bridges/a2a" + clienta2a "github.com/agent-dance/agent-adaptor/pkg/clients/a2a" +) + +type eventMapper struct { + base DelegationEvent + started bool + openMessage string +} + +func newEventMapper(base DelegationEvent) *eventMapper { + base.Protocol = ProtocolA2A + return &eventMapper{base: base} +} + +func (m *eventMapper) Started(taskID, contextID string) DelegationEvent { + m.started = true + ev := m.base + ev.Kind = DelegationStarted + ev.RemoteTaskID = taskID + ev.RemoteContextID = contextID + ev.Status = "started" + return ev +} + +func (m *eventMapper) Map(event clienta2a.Event) []DelegationEvent { + out := []DelegationEvent{} + if !m.started { + out = append(out, m.Started(event.TaskID, event.ContextID)) + } + switch event.Kind { + case clienta2a.EventStatus: + out = append(out, m.statusEvent(event.TaskID, event.ContextID, event.Status)) + case clienta2a.EventTask: + if event.Task != nil { + out = append(out, m.taskEvents(*event.Task)...) + } + case clienta2a.EventMessage: + if event.Message != nil { + out = append(out, m.messageEvents(*event.Message)...) + } + case clienta2a.EventArtifact: + if event.Artifact != nil { + out = append(out, m.artifactEvent(event.TaskID, event.ContextID, *event.Artifact)) + } + case clienta2a.EventTerminal: + out = append(out, m.terminalEvents(event)...) + } + return out +} + +func (m *eventMapper) taskEvents(task clienta2a.Task) []DelegationEvent { + out := []DelegationEvent{} + if task.Status.State != "" { + out = append(out, m.statusEvent(task.ID, task.ContextID, &task.Status)) + } + for _, artifact := range task.Artifacts { + out = append(out, m.artifactEvent(task.ID, task.ContextID, artifact)) + } + if executionFinalState(task.Status.State) { + out = append(out, m.terminalForState(task.ID, task.ContextID, task.Status.State, task.Raw)) + } + return out +} + +func (m *eventMapper) statusEvent(taskID, contextID string, status *clienta2a.TaskStatus) DelegationEvent { + ev := m.base + ev.Kind = DelegationStatus + ev.RemoteTaskID = taskID + ev.RemoteContextID = contextID + if status != nil { + ev.Status = string(status.State) + ev.Raw = map[string]any{"status": string(status.State)} + if status.Message != nil { + ev.Text = textFromMessage(*status.Message) + } + } + return ev +} + +func (m *eventMapper) messageEvents(msg clienta2a.Message) []DelegationEvent { + text := textFromMessage(msg) + if text == "" { + return nil + } + messageID := msg.ID + if messageID == "" { + messageID = m.base.DelegationID + ":message" + } + out := []DelegationEvent{} + if m.openMessage != messageID { + if m.openMessage != "" { + end := m.base + end.Kind = DelegationTextEnd + end.RemoteMessageID = m.openMessage + out = append(out, end) + } + start := m.base + start.Kind = DelegationTextStart + start.RemoteTaskID = msg.TaskID + start.RemoteContextID = msg.ContextID + start.RemoteMessageID = messageID + out = append(out, start) + m.openMessage = messageID + } + delta := m.base + delta.Kind = DelegationTextDelta + delta.RemoteTaskID = msg.TaskID + delta.RemoteContextID = msg.ContextID + delta.RemoteMessageID = messageID + delta.Delta = text + out = append(out, delta) + return out +} + +func (m *eventMapper) artifactEvent(taskID, contextID string, artifact clienta2a.Artifact) DelegationEvent { + if artifact.Name == bridgea2a.ArtifactAssistantOutput { + ev := m.base + ev.Kind = DelegationTextDelta + ev.RemoteTaskID = taskID + ev.RemoteContextID = contextID + ev.RemoteArtifactID = artifact.ID + ev.RemoteMessageID = bridgea2a.ArtifactAssistantOutput + ev.Delta = textFromParts(artifact.Parts) + return ev + } + ev := m.base + ev.Kind = DelegationArtifactCreated + ev.RemoteTaskID = taskID + ev.RemoteContextID = contextID + ev.RemoteArtifactID = artifact.ID + ev.Artifact = &DelegationArtifact{ + ID: artifact.ID, + Name: artifact.Name, + Description: artifact.Description, + URI: firstURI(artifact.Parts), + MediaType: firstMediaType(artifact.Parts), + Metadata: cloneAnyMap(artifact.Metadata), + } + ev.Raw = artifact.Raw + return ev +} + +func (m *eventMapper) terminalEvents(event clienta2a.Event) []DelegationEvent { + if event.Task != nil { + return []DelegationEvent{m.terminalForState(event.Task.ID, event.Task.ContextID, event.Task.Status.State, event.Task.Raw)} + } + if event.Status != nil { + return []DelegationEvent{m.terminalForState(event.TaskID, event.ContextID, event.Status.State, event.Raw)} + } + if event.Message != nil { + out := m.messageEvents(*event.Message) + out = append(out, m.terminalForState(event.TaskID, event.ContextID, clienta2a.TaskStateCompleted, event.Raw)) + return out + } + return []DelegationEvent{m.terminalForState(event.TaskID, event.ContextID, clienta2a.TaskStateCompleted, event.Raw)} +} + +func (m *eventMapper) terminalForState(taskID, contextID string, state clienta2a.TaskState, raw map[string]any) DelegationEvent { + ev := m.base + ev.RemoteTaskID = taskID + ev.RemoteContextID = contextID + ev.Status = string(state) + ev.Raw = raw + switch state { + case clienta2a.TaskStateCompleted: + ev.Kind = DelegationFinished + case clienta2a.TaskStateCanceled: + ev.Kind = DelegationCancelled + case clienta2a.TaskStateInputRequired: + ev.Kind = DelegationInputRequired + default: + ev.Kind = DelegationFailed + ev.Error = &DelegationError{Code: "remote_failed", Message: "remote task failed", RemoteStatus: string(state)} + } + return ev +} + +func resultFromTask(base DelegationResult, task clienta2a.Task) DelegationResult { + base.RemoteTaskID = task.ID + base.RemoteContextID = task.ContextID + base.Status = statusFromState(task.Status.State) + base.RawTask = map[string]any{"provider": ProtocolA2A, "task_id": task.ID} + for _, msg := range task.Messages { + text := textFromMessage(msg) + if text != "" { + base.Messages = append(base.Messages, DelegationMessage{Role: msg.Role, Text: text}) + if base.Summary == "" { + base.Summary = text + } + } + } + for _, artifact := range task.Artifacts { + if artifact.Name == bridgea2a.ArtifactAgentAdaptorResult { + applyAgentAdaptorResult(&base, artifact) + continue + } + if artifact.Name == bridgea2a.ArtifactAssistantOutput { + if text := textFromParts(artifact.Parts); text != "" && base.Summary == "" { + base.Summary = text + } + continue + } + base.Artifacts = append(base.Artifacts, DelegationArtifact{ID: artifact.ID, Name: artifact.Name, Description: artifact.Description, URI: firstURI(artifact.Parts), MediaType: firstMediaType(artifact.Parts), Metadata: cloneAnyMap(artifact.Metadata)}) + } + if base.Status != "completed" && base.Error == nil { + base.Error = &DelegationError{Code: errorCodeFromState(task.Status.State), Message: "remote task did not complete successfully", RemoteStatus: string(task.Status.State)} + } + return base +} + +func applyAgentAdaptorResult(result *DelegationResult, artifact clienta2a.Artifact) { + for _, part := range artifact.Parts { + if part.Kind != clienta2a.PartData { + continue + } + raw, err := json.Marshal(part.Data) + if err != nil { + continue + } + var payload struct { + Summary string `json:"summary"` + Output string `json:"output"` + Result map[string]any `json:"result"` + } + if err := json.Unmarshal(raw, &payload); err != nil { + continue + } + if payload.Summary != "" { + result.Summary = payload.Summary + } else if payload.Output != "" { + result.Summary = payload.Output + } + if len(payload.Result) > 0 { + result.Metadata = payload.Result + } + } +} + +func executionFinalState(state clienta2a.TaskState) bool { + return state.Terminal() || state == clienta2a.TaskStateInputRequired +} + +func statusFromState(state clienta2a.TaskState) string { + switch state { + case clienta2a.TaskStateCompleted: + return "completed" + case clienta2a.TaskStateCanceled: + return "cancelled" + case clienta2a.TaskStateInputRequired: + return "input_required" + default: + if state == "" { + return "unknown" + } + return "failed" + } +} + +func errorCodeFromState(state clienta2a.TaskState) string { + switch state { + case clienta2a.TaskStateCanceled: + return "remote_cancelled" + case clienta2a.TaskStateInputRequired: + return "input_required" + case clienta2a.TaskStateRejected: + return "remote_rejected" + default: + return "remote_failed" + } +} + +func textFromMessage(msg clienta2a.Message) string { + return textFromParts(msg.Parts) +} + +func textFromParts(parts []clienta2a.Part) string { + var b strings.Builder + for _, part := range parts { + if part.Kind == clienta2a.PartText && part.Text != "" { + b.WriteString(part.Text) + } + } + return b.String() +} + +func firstURI(parts []clienta2a.Part) string { + for _, part := range parts { + if part.URL != "" { + return part.URL + } + } + return "" +} + +func firstMediaType(parts []clienta2a.Part) string { + for _, part := range parts { + if part.MediaType != "" { + return part.MediaType + } + } + return "" +} diff --git a/pkg/hosttools/a2adelegation/mcpserver.go b/pkg/hosttools/a2adelegation/mcpserver.go new file mode 100644 index 0000000..953b4f5 --- /dev/null +++ b/pkg/hosttools/a2adelegation/mcpserver.go @@ -0,0 +1,273 @@ +package a2adelegation + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "time" +) + +type ToolInput struct { + Agent string `json:"agent"` + Objective string `json:"objective"` + Input ToolInputBody `json:"input,omitempty"` + Constraints ToolConstraints `json:"constraints,omitempty"` +} + +type ToolInputBody struct { + Prompt string `json:"prompt,omitempty"` + Context string `json:"context,omitempty"` + Artifacts []InputArtifact `json:"artifacts,omitempty"` +} + +type ToolConstraints struct { + TimeoutSeconds int `json:"timeout_seconds,omitempty"` + Stream bool `json:"stream,omitempty"` + MaxArtifacts int `json:"max_artifacts,omitempty"` + MaxArtifactsSet bool `json:"-"` +} + +func ToolSchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "agent": map[string]any{"type": "string", "description": "Host-curated remote agent registry key."}, + "objective": map[string]any{"type": "string", "description": "Bounded task objective for the remote agent."}, + "input": map[string]any{ + "type": "object", + "properties": map[string]any{ + "prompt": map[string]any{"type": "string"}, + "context": map[string]any{"type": "string"}, + "artifacts": map[string]any{ + "type": "array", + "items": map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string"}, + "uri": map[string]any{"type": "string"}, + "mime_type": map[string]any{"type": "string"}, + }, + "additionalProperties": false, + }, + }, + }, + "additionalProperties": false, + }, + "constraints": map[string]any{ + "type": "object", + "properties": map[string]any{ + "timeout_seconds": map[string]any{"type": "integer", "minimum": 1}, + "stream": map[string]any{"type": "boolean"}, + "max_artifacts": map[string]any{"type": "integer", "minimum": 0}, + }, + "additionalProperties": false, + }, + }, + "required": []string{"agent", "objective"}, + "additionalProperties": false, + } +} + +func ParseToolInput(raw []byte) (ToolInput, error) { + var envelope map[string]json.RawMessage + if err := json.Unmarshal(raw, &envelope); err == nil { + if _, ok := envelope["endpoint_url"]; ok { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "endpoint_url is not allowed; use a host-curated agent key"} + } + for key := range envelope { + switch key { + case "agent", "objective", "input", "constraints": + default: + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "unknown field " + key} + } + } + } + var input ToolInput + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.DisallowUnknownFields() + if err := dec.Decode(&input); err != nil { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: err.Error()} + } + if input.Agent == "" { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "agent is required"} + } + if input.Objective == "" { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "objective is required"} + } + if rawConstraints, ok := envelope["constraints"]; ok { + var constraints map[string]json.RawMessage + if err := json.Unmarshal(rawConstraints, &constraints); err == nil { + if _, ok := constraints["max_artifacts"]; ok { + input.Constraints.MaxArtifactsSet = true + } + } + } + if input.Constraints.MaxArtifactsSet && input.Constraints.MaxArtifacts < 0 { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "max_artifacts must be non-negative"} + } + return input, nil +} + +type MCPServerOptions struct { + RunID string + ParentToolCallID string + Tenant string + BearerToken string +} + +type MCPServer struct { + Delegator *Delegator + Options MCPServerOptions +} + +func NewMCPServer(delegator *Delegator, opts MCPServerOptions) *MCPServer { + return &MCPServer{Delegator: delegator, Options: opts} +} + +func (s *MCPServer) Handler() http.Handler { + return http.HandlerFunc(s.ServeHTTP) +} + +func (s *MCPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if err := s.authorize(r); err != nil { + http.Error(w, err.Error(), http.StatusUnauthorized) + return + } + defer r.Body.Close() + var req rpcRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeRPC(w, rpcError(nil, -32700, "parse error: "+err.Error())) + return + } + if len(req.ID) == 0 && strings.HasPrefix(req.Method, "notifications/") { + w.WriteHeader(http.StatusNoContent) + return + } + writeRPC(w, s.handle(r, req)) +} + +func (s *MCPServer) authorize(r *http.Request) error { + if s == nil || s.Options.BearerToken == "" { + return nil + } + got := strings.TrimSpace(r.Header.Get("Authorization")) + want := "Bearer " + s.Options.BearerToken + if got != want { + return errors.New("unauthorized") + } + return nil +} + +func (s *MCPServer) handle(r *http.Request, req rpcRequest) map[string]any { + switch req.Method { + case "initialize": + return result(req.ID, map[string]any{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]any{"tools": map[string]any{}}, + "serverInfo": map[string]any{"name": "agent-adaptor-a2a-delegation", "version": "0.1.0"}, + }) + case "ping": + return result(req.ID, map[string]any{}) + case "tools/list": + return result(req.ID, map[string]any{"tools": []map[string]any{{ + "name": DelegateToolName, + "description": "Delegate a bounded task to a host-curated remote A2A agent and return a structured final result.", + "inputSchema": ToolSchema(), + }}}) + case "tools/call": + return s.handleToolCall(r, req) + default: + return rpcError(req.ID, -32601, "method not found") + } +} + +func (s *MCPServer) handleToolCall(r *http.Request, req rpcRequest) map[string]any { + if s == nil || s.Delegator == nil { + return rpcError(req.ID, -32000, "delegator is not configured") + } + var params struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + } + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + return rpcError(req.ID, -32602, "invalid tools/call params: "+err.Error()) + } + if params.Name != DelegateToolName { + return rpcError(req.ID, -32602, "unknown tool "+params.Name) + } + input, err := ParseToolInput(params.Arguments) + if err != nil { + return toolResult(req.ID, DelegationResult{Status: "failed", Error: delegationErr(err)}, true) + } + request := DelegationRequest{ + RunID: s.Options.RunID, + ParentToolCallID: s.Options.ParentToolCallID, + Agent: input.Agent, + Objective: input.Objective, + Prompt: input.Input.Prompt, + Context: input.Input.Context, + Artifacts: append([]InputArtifact(nil), input.Input.Artifacts...), + Stream: input.Constraints.Stream, + Tenant: s.Options.Tenant, + } + if input.Constraints.MaxArtifactsSet { + request.MaxArtifacts = &input.Constraints.MaxArtifacts + } + if input.Constraints.TimeoutSeconds > 0 { + request.Timeout = time.Duration(input.Constraints.TimeoutSeconds) * time.Second + } + out, err := s.Delegator.Delegate(r.Context(), request) + return toolResult(req.ID, out, err != nil) +} + +func delegationErr(err error) *DelegationError { + var derr *DelegationError + if errors.As(err, &derr) { + return derr + } + return &DelegationError{Code: "delegation_error", Message: err.Error()} +} + +type rpcRequest struct { + JSONRPC string `json:"jsonrpc,omitempty"` + ID json.RawMessage `json:"id,omitempty"` + Method string `json:"method,omitempty"` + Params json.RawMessage `json:"params,omitempty"` +} + +func toolResult(id json.RawMessage, value DelegationResult, isError bool) map[string]any { + raw, err := json.Marshal(value) + if err != nil { + return rpcError(id, -32603, "marshal delegation result: "+err.Error()) + } + return result(id, map[string]any{ + "content": []map[string]any{{"type": "text", "text": string(raw)}}, + "isError": isError, + }) +} + +func result(id json.RawMessage, value any) map[string]any { + return map[string]any{"jsonrpc": "2.0", "id": json.RawMessage(id), "result": value} +} + +func rpcError(id json.RawMessage, code int, message string) map[string]any { + return map[string]any{"jsonrpc": "2.0", "id": json.RawMessage(id), "error": map[string]any{"code": code, "message": message}} +} + +func writeRPC(w http.ResponseWriter, payload map[string]any) { + w.Header().Set("Content-Type", "application/json") + if payload == nil { + w.WriteHeader(http.StatusNoContent) + return + } + if err := json.NewEncoder(w).Encode(payload); err != nil { + http.Error(w, fmt.Sprintf("encode rpc response: %v", err), http.StatusInternalServerError) + } +} diff --git a/pkg/hosttools/a2adelegation/types.go b/pkg/hosttools/a2adelegation/types.go new file mode 100644 index 0000000..061552b --- /dev/null +++ b/pkg/hosttools/a2adelegation/types.go @@ -0,0 +1,238 @@ +// Package a2adelegation contains host-owned tools for exposing curated remote +// A2A agents as visual subagents of a local parent run. +// +// The package deliberately stays above the core SDK: it consumes +// pkg/clients/a2a DTOs, emits UI-facing delegation events, and leaves concrete +// local adapters unaware of A2A. +package a2adelegation + +import ( + "fmt" + "net/http" + "strings" + "time" + + clienta2a "github.com/agent-dance/agent-adaptor/pkg/clients/a2a" +) + +const ProtocolA2A = "a2a" + +const DelegateToolName = "delegate_to_agent" + +type DelegationEventKind string + +const ( + DelegationStarted DelegationEventKind = "subagent.started" + DelegationStatus DelegationEventKind = "subagent.status" + DelegationTextStart DelegationEventKind = "subagent.text.start" + DelegationTextDelta DelegationEventKind = "subagent.text.delta" + DelegationTextEnd DelegationEventKind = "subagent.text.end" + DelegationArtifactCreated DelegationEventKind = "subagent.artifact" + DelegationInputRequired DelegationEventKind = "subagent.input_required" + DelegationFinished DelegationEventKind = "subagent.finished" + DelegationFailed DelegationEventKind = "subagent.failed" + DelegationCancelled DelegationEventKind = "subagent.cancelled" +) + +type RemoteAgentSpec struct { + Key string + DisplayName string + Protocol string + AgentCardURL string + AgentCard *clienta2a.AgentCard + Tenant string + Auth clienta2a.Auth + HTTPClient *http.Client + TrustedAuthOrigins []string + AcceptedOutputModes []string + PreferredTransports []clienta2a.TransportProtocol + Policy DelegationPolicy +} + +type DelegationPolicy struct { + MaxTimeout time.Duration + AllowInputRequired bool + RequireStreaming bool + PollInterval time.Duration + MaxPolls int + MaxArtifactBytes int64 +} + +type Registry struct { + agents map[string]RemoteAgentSpec +} + +func NewRegistry(specs ...RemoteAgentSpec) (*Registry, error) { + r := &Registry{agents: map[string]RemoteAgentSpec{}} + for _, spec := range specs { + if err := r.Register(spec); err != nil { + return nil, err + } + } + return r, nil +} + +func (r *Registry) Register(spec RemoteAgentSpec) error { + if r.agents == nil { + r.agents = map[string]RemoteAgentSpec{} + } + key := strings.TrimSpace(spec.Key) + if key == "" { + return &DelegationError{Code: "invalid_agent", Message: "remote agent key is required"} + } + if _, exists := r.agents[key]; exists { + return &DelegationError{Code: "duplicate_agent", Message: fmt.Sprintf("remote agent %q already registered", key)} + } + spec.Key = key + if spec.Protocol == "" { + spec.Protocol = ProtocolA2A + } + if spec.Protocol != ProtocolA2A { + return &DelegationError{Code: "unsupported_protocol", Message: fmt.Sprintf("remote agent %q uses unsupported protocol %q", key, spec.Protocol)} + } + if strings.TrimSpace(spec.AgentCardURL) == "" && spec.AgentCard == nil { + return &DelegationError{Code: "invalid_agent", Message: fmt.Sprintf("remote agent %q requires a host-curated AgentCardURL or AgentCard", key)} + } + r.agents[key] = cloneRemoteAgentSpec(spec) + return nil +} + +func (r *Registry) Lookup(key string) (RemoteAgentSpec, bool) { + if r == nil || len(r.agents) == 0 { + return RemoteAgentSpec{}, false + } + spec, ok := r.agents[strings.TrimSpace(key)] + return cloneRemoteAgentSpec(spec), ok +} + +func (r *Registry) Keys() []string { + if r == nil || len(r.agents) == 0 { + return nil + } + keys := make([]string, 0, len(r.agents)) + for key := range r.agents { + keys = append(keys, key) + } + return keys +} + +type DelegationRequest struct { + RunID string + ParentToolCallID string + Agent string + Objective string + Prompt string + Context string + Artifacts []InputArtifact + MaxArtifacts *int + Timeout time.Duration + Stream bool + Tenant string + Metadata map[string]any +} + +type InputArtifact struct { + Name string `json:"name,omitempty"` + URI string `json:"uri,omitempty"` + MediaType string `json:"mime_type,omitempty"` +} + +type DelegationArtifact struct { + ID string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Description string `json:"description,omitempty"` + URI string `json:"uri,omitempty"` + MediaType string `json:"mime_type,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +type DelegationEvent struct { + RunID string + ParentToolCallID string + DelegationID string + AgentKey string + AgentName string + Protocol string + + RemoteTaskID string + RemoteContextID string + RemoteMessageID string + RemoteArtifactID string + + Kind DelegationEventKind + Delta string + Text string + Artifact *DelegationArtifact + Status string + Error *DelegationError + Raw map[string]any + Time time.Time +} + +type DelegationResult struct { + DelegationID string `json:"delegation_id"` + Agent string `json:"agent"` + RemoteProtocol string `json:"remote_protocol"` + RemoteTaskID string `json:"remote_task_id,omitempty"` + RemoteContextID string `json:"remote_context_id,omitempty"` + Status string `json:"status"` + Summary string `json:"summary,omitempty"` + Artifacts []DelegationArtifact `json:"artifacts,omitempty"` + Messages []DelegationMessage `json:"messages,omitempty"` + Error *DelegationError `json:"error,omitempty"` + RawTask map[string]any `json:"raw_task,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` +} + +type DelegationMessage struct { + Role string `json:"role,omitempty"` + Text string `json:"text,omitempty"` +} + +type DelegationError struct { + Code string `json:"code"` + Message string `json:"message"` + Retryable bool `json:"retryable,omitempty"` + RemoteStatus string `json:"remote_status,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +func (e *DelegationError) Error() string { + if e == nil { + return "" + } + if e.Message != "" { + return e.Message + } + if e.Code != "" { + return e.Code + } + return "delegation failed" +} + +func cloneRemoteAgentSpec(spec RemoteAgentSpec) RemoteAgentSpec { + spec.TrustedAuthOrigins = append([]string(nil), spec.TrustedAuthOrigins...) + spec.AcceptedOutputModes = append([]string(nil), spec.AcceptedOutputModes...) + spec.PreferredTransports = append([]clienta2a.TransportProtocol(nil), spec.PreferredTransports...) + if spec.AgentCard != nil { + card := *spec.AgentCard + card.DefaultInputModes = append([]string(nil), card.DefaultInputModes...) + card.DefaultOutputModes = append([]string(nil), card.DefaultOutputModes...) + card.Skills = append([]clienta2a.Skill(nil), card.Skills...) + card.SupportedInterfaces = append([]clienta2a.AgentInterface(nil), card.SupportedInterfaces...) + card.Raw = cloneAnyMap(card.Raw) + spec.AgentCard = &card + } + return spec +} + +func cloneAnyMap(in map[string]any) map[string]any { + if len(in) == 0 { + return nil + } + out := make(map[string]any, len(in)) + for k, v := range in { + out[k] = v + } + return out +} diff --git a/runner.go b/runner.go index 3611814..2864803 100644 --- a/runner.go +++ b/runner.go @@ -309,10 +309,6 @@ func (r *runnerImpl) resolveInvocation(ctx context.Context, prompt string, opts defaultRefs := r.sdk.selectedRefsFor(r.name, defaults.Skills) runRefs := cloneSkillRefs(resolvedOpts.skills) - mcpPayload, err := resolveMCPPayload(defaults.MCP, resolvedOpts.mcp, r.binding.Adapter().Descriptor().MCP) - if err != nil { - return resolvedInvocation{}, nil, err - } policy, err := mergeRunPolicy(defaults.RunPolicy, resolvedOpts.runPolicy) if err != nil { @@ -392,6 +388,12 @@ func (r *runnerImpl) resolveInvocation(ctx context.Context, prompt string, opts _ = r.sdk.workspaceManager.Release(releaseCtx, workspace, WorkspaceReleaseKeep) } + mcpPayload, err := resolveMCPPayloadWithRuntime(defaults.MCP, resolvedOpts.mcp, runtimePayload.Ensured, r.binding.Adapter().Descriptor().MCP) + if err != nil { + cleanup() + return resolvedInvocation{}, nil, err + } + skillPayload, _, _, err := r.sdk.resolveSkills(ctx, identity, defaultRefs, runRefs, defaults.Skills) if err != nil { cleanup() diff --git a/runtime_admin_test.go b/runtime_admin_test.go index d3f9eee..78a8e48 100644 --- a/runtime_admin_test.go +++ b/runtime_admin_test.go @@ -8,7 +8,8 @@ import ( ) type runtimeAdminDriver struct { - lastReq agentadaptor.DriverRunRequest + lastReq agentadaptor.DriverRunRequest + mcpCapability agentadaptor.MCPCapability } func (d *runtimeAdminDriver) Descriptor() agentadaptor.DriverDescriptor { @@ -21,6 +22,7 @@ func (d *runtimeAdminDriver) Descriptor() agentadaptor.DriverDescriptor { }, }, Runtime: agentadaptor.RuntimeCapability{ReportsServices: true}, + MCP: d.mcpCapability, } } @@ -189,6 +191,167 @@ func TestRuntimeServicesFlowThroughSingleExecutionPath(t *testing.T) { } } +func TestRuntimeServiceMetadataInjectsMCPIntoDriverAndProfile(t *testing.T) { + driver := &runtimeAdminDriver{ + mcpCapability: agentadaptor.MCPCapability{Supported: true, HTTP: true}, + } + runtimeManager := &runtimeMCPManager{ + ref: agentadaptor.RuntimeServiceRef{ + ID: "svc-delegation", + Name: "a2a-delegation", + URL: "http://127.0.0.1:43127/mcp", + Metadata: map[string]string{ + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.key": "delegate-a2a", + "agentadaptor.mcp.transport": "http", + "agentadaptor.mcp.headers_json": `{"X-Run-Token":"env:DELEGATION_TOKEN"}`, + "agentadaptor.mcp.bearer_token_env_var": "DELEGATION_TOKEN", + "agentadaptor.mcp.required": "true", + "agentadaptor.mcp.required_reason": "visual A2A subagent delegation", + "delegation.secret_should_not_be_promoted": "sk-test-secret", + }, + }, + } + sdk := agentadaptor.New( + agentadaptor.WithDefaultAgent(agentadaptor.Bind(driver, fakeConfig{Label: "runtime"}, + agentadaptor.WithDefaultMCP(agentadaptor.MCPConfig{Servers: []agentadaptor.MCPServerSpec{{Key: "host", Transport: agentadaptor.MCPTransportHTTP, URL: "http://127.0.0.1:1/mcp"}}}), + agentadaptor.WithDefaultRuntimeServices(agentadaptor.RuntimeServiceSpec{Name: "a2a-delegation"}), + )), + agentadaptor.WithRuntimeServiceManager(runtimeManager), + ) + + if _, err := sdk.Run(context.Background(), "use delegation"); err != nil { + t.Fatalf("run: %v", err) + } + if len(driver.lastReq.MCP.Servers) != 2 { + t.Fatalf("expected host + runtime MCP servers, got %#v", driver.lastReq.MCP.Servers) + } + var injected agentadaptor.MCPServerSpec + for _, server := range driver.lastReq.MCP.Servers { + if server.Key == "delegate-a2a" { + injected = server + } + } + if injected.Key == "" { + t.Fatalf("missing runtime-injected MCP server: %#v", driver.lastReq.MCP.Servers) + } + if injected.Transport != agentadaptor.MCPTransportHTTP || injected.URL != "http://127.0.0.1:43127/mcp" { + t.Fatalf("unexpected injected server endpoint: %#v", injected) + } + if injected.Headers["X-Run-Token"] != "env:DELEGATION_TOKEN" || injected.BearerTokenEnvVar != "DELEGATION_TOKEN" { + t.Fatalf("expected run-scoped auth references, got %#v", injected) + } + if !injected.Required || injected.RequiredReason != "visual A2A subagent delegation" { + t.Fatalf("expected required runtime MCP server, got %#v", injected) + } + if driver.lastReq.ProfilePayload.MCP.Fingerprint != driver.lastReq.MCP.Fingerprint { + t.Fatalf("profile MCP fingerprint %q did not match driver MCP fingerprint %q", driver.lastReq.ProfilePayload.MCP.Fingerprint, driver.lastReq.MCP.Fingerprint) + } + if driver.lastReq.ProfilePayload.Fingerprint == "" { + t.Fatal("expected profile fingerprint to include runtime MCP payload") + } +} + +func TestRuntimeServiceMCPMetadataRejectsMalformedJSON(t *testing.T) { + driver := &runtimeAdminDriver{ + mcpCapability: agentadaptor.MCPCapability{Supported: true, HTTP: true}, + } + sdk := agentadaptor.New( + agentadaptor.WithDefaultAgent(agentadaptor.Bind(driver, fakeConfig{Label: "runtime"}, + agentadaptor.WithDefaultRuntimeServices(agentadaptor.RuntimeServiceSpec{Name: "bad-mcp"}), + )), + agentadaptor.WithRuntimeServiceManager(&runtimeMCPManager{ref: agentadaptor.RuntimeServiceRef{ + Name: "bad-mcp", + URL: "http://127.0.0.1:43127/mcp", + Metadata: map[string]string{ + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.headers_json": `{"Authorization": 123}`, + }, + }}), + ) + + if _, err := sdk.Run(context.Background(), "use delegation"); err == nil { + t.Fatal("expected malformed runtime MCP metadata to fail") + } +} + +func TestRuntimeServiceMCPDuplicateKeyIsRejected(t *testing.T) { + driver := &runtimeAdminDriver{ + mcpCapability: agentadaptor.MCPCapability{Supported: true, HTTP: true}, + } + sdk := agentadaptor.New( + agentadaptor.WithDefaultAgent(agentadaptor.Bind(driver, fakeConfig{Label: "runtime"}, + agentadaptor.WithDefaultMCP(agentadaptor.MCPConfig{Servers: []agentadaptor.MCPServerSpec{{Key: "delegate-a2a", Transport: agentadaptor.MCPTransportHTTP, URL: "http://127.0.0.1:1/mcp"}}}), + agentadaptor.WithDefaultRuntimeServices(agentadaptor.RuntimeServiceSpec{Name: "a2a-delegation"}), + )), + agentadaptor.WithRuntimeServiceManager(&runtimeMCPManager{ref: agentadaptor.RuntimeServiceRef{ + Name: "a2a-delegation", + URL: "http://127.0.0.1:43127/mcp", + Metadata: map[string]string{ + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.key": "delegate-a2a", + }, + }}), + ) + + if _, err := sdk.Run(context.Background(), "use delegation"); err == nil { + t.Fatal("expected duplicate runtime MCP key to fail") + } +} + +func TestRuntimeMCPChangesSessionFingerprint(t *testing.T) { + driver := &runtimeAdminDriver{ + mcpCapability: agentadaptor.MCPCapability{Supported: true, HTTP: true}, + } + runtimeManager := &runtimeMCPManager{} + sdk := agentadaptor.New( + agentadaptor.WithDefaultAgent(agentadaptor.Bind(driver, fakeConfig{Label: "runtime"}, + agentadaptor.WithDefaultRuntimeServices(agentadaptor.RuntimeServiceSpec{Name: "delegate"}), + )), + agentadaptor.WithRuntimeServiceManager(runtimeManager), + ) + + runtimeManager.ref = agentadaptor.RuntimeServiceRef{ + Name: "delegate", + URL: "http://127.0.0.1:43127/mcp", + Metadata: map[string]string{ + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.key": "delegate-a", + }, + } + if _, err := sdk.Run(context.Background(), "first"); err != nil { + t.Fatalf("first run: %v", err) + } + first := driver.lastReq.ProfilePayload.Fingerprint + + runtimeManager.ref = agentadaptor.RuntimeServiceRef{ + Name: "delegate", + URL: "http://127.0.0.1:43128/mcp", + Metadata: map[string]string{ + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.key": "delegate-b", + }, + } + if _, err := sdk.Run(context.Background(), "second"); err != nil { + t.Fatalf("second run: %v", err) + } + if driver.lastReq.ProfilePayload.Fingerprint == first { + t.Fatal("expected runtime MCP profile fingerprint to change") + } +} + +type runtimeMCPManager struct { + ref agentadaptor.RuntimeServiceRef +} + +func (m *runtimeMCPManager) Ensure(_ context.Context, _ agentadaptor.RuntimeServiceRequest) ([]agentadaptor.RuntimeServiceRef, error) { + return []agentadaptor.RuntimeServiceRef{m.ref}, nil +} + +func (m *runtimeMCPManager) ReleaseByRun(_ context.Context, _ string) error { return nil } + +func (m *runtimeMCPManager) ReleaseByLabels(_ context.Context, _ map[string]string) error { return nil } + func TestAdminExposesConfigSchemaQuotaAndDetectedModel(t *testing.T) { driver := &runtimeAdminDriver{} sdk := agentadaptor.New( From 3f03f2fa376362a46f68b88c06ba376e0cec0e8f Mon Sep 17 00:00:00 2001 From: buthim Date: Tue, 7 Jul 2026 15:41:04 +0800 Subject: [PATCH 2/6] fix: address delegation runtime review issues Co-Authored-By: Claude Opus 4.8 --- docs/streaming.md | 5 +- pkg/bridges/subagentstream/mux.go | 18 +- pkg/bridges/subagentstream/mux_test.go | 75 ++++++ pkg/hosttools/a2adelegation/delegator.go | 122 +++++++-- pkg/hosttools/a2adelegation/delegator_test.go | 254 ++++++++++++++++-- pkg/hosttools/a2adelegation/events.go | 21 +- pkg/hosttools/a2adelegation/mapping.go | 7 +- 7 files changed, 454 insertions(+), 48 deletions(-) diff --git a/docs/streaming.md b/docs/streaming.md index d449543..cb41e6d 100644 --- a/docs/streaming.md +++ b/docs/streaming.md @@ -167,7 +167,10 @@ mux.Handle("/v1/chat", sse.Handler(sdk, sse.Options{ // The MCP tool server for the same run calls delegator.Delegate(...), which // returns only the final structured DelegationResult to the parent model while -// publishing live progress onto bus.SubscribeRun(ctx, handle.RunID()). +// publishing live progress onto bus.SubscribeRun(ctx, handle.RunID()). When the +// run is no longer observable, call bus.ClearRun(handle.RunID()) (or equivalent +// host lifecycle cleanup) so replay and terminal-dedupe state do not grow without +// bound in long-lived services. ``` `SubagentBus` is honored by the stock SSE handler only when `Protocol: sse.AGUI`. diff --git a/pkg/bridges/subagentstream/mux.go b/pkg/bridges/subagentstream/mux.go index bdae351..33b2698 100644 --- a/pkg/bridges/subagentstream/mux.go +++ b/pkg/bridges/subagentstream/mux.go @@ -90,12 +90,24 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < defer close(out) parent := agui.WrapWithContext(ctx, handle) var subagents <-chan a2adelegation.DelegationEvent + var cancelSubagents context.CancelFunc if opts.Bus != nil { - subagents = opts.Bus.SubscribeRun(ctx, handle.RunID()) + subagentCtx, cancel := context.WithCancel(ctx) + cancelSubagents = cancel + defer cancelSubagents() + subagents = opts.Bus.SubscribeRun(subagentCtx, handle.RunID()) + } + stopSubagents := func() { + if cancelSubagents != nil { + cancelSubagents() + cancelSubagents = nil + } + subagents = nil } for parent != nil || subagents != nil { select { case <-ctx.Done(): + stopSubagents() cancelHandle(handle) return case ev, ok := <-parent: @@ -104,7 +116,7 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < if !drainSubagents(subagents) { return } - subagents = nil + stopSubagents() continue } if isTerminalAGUI(ev) { @@ -115,7 +127,7 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < return } parent = nil - subagents = nil + stopSubagents() continue } if !sendAGUI(ev) { diff --git a/pkg/bridges/subagentstream/mux_test.go b/pkg/bridges/subagentstream/mux_test.go index 9a2eafc..4eb1d6e 100644 --- a/pkg/bridges/subagentstream/mux_test.go +++ b/pkg/bridges/subagentstream/mux_test.go @@ -170,6 +170,43 @@ func TestWrapCancelUsesBestEffortContext(t *testing.T) { } } +func TestWrapCancelsSubagentSubscriptionWhenParentTerminates(t *testing.T) { + t.Parallel() + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload, 1), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := newTrackingBus() + out := subagentstream.Wrap(context.Background(), handle, subagentstream.MuxOptions{Bus: bus}) + handle.stream <- agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunFinished, ThreadID: "thread-1", RunID: "run-1"} + close(handle.stream) + close(handle.done) + for range out { + } + bus.assertCanceled(t, "run-1") +} + +func TestWrapCancelsSubagentSubscriptionWhenParentStreamCloses(t *testing.T) { + t.Parallel() + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := newTrackingBus() + out := subagentstream.Wrap(context.Background(), handle, subagentstream.MuxOptions{Bus: bus}) + close(handle.stream) + close(handle.done) + for range out { + } + bus.assertCanceled(t, "run-1") +} + func collectMuxEvents(t *testing.T, ch <-chan subagentstream.Event, want int) []subagentstream.Event { t.Helper() out := make([]subagentstream.Event, 0, want) @@ -187,6 +224,44 @@ func collectMuxEvents(t *testing.T, ch <-chan subagentstream.Event, want int) [] return out } +type trackingBus struct { + mu sync.Mutex + runID string + ctxDone <-chan struct{} + events chan a2adelegation.DelegationEvent +} + +func newTrackingBus() *trackingBus { + return &trackingBus{events: make(chan a2adelegation.DelegationEvent)} +} + +func (b *trackingBus) SubscribeRun(ctx context.Context, runID string) <-chan a2adelegation.DelegationEvent { + b.mu.Lock() + b.runID = runID + b.ctxDone = ctx.Done() + b.mu.Unlock() + return b.events +} + +func (b *trackingBus) assertCanceled(t *testing.T, wantRunID string) { + t.Helper() + b.mu.Lock() + runID := b.runID + done := b.ctxDone + b.mu.Unlock() + if runID != wantRunID { + t.Fatalf("subscribed runID: got %q want %q", runID, wantRunID) + } + if done == nil { + t.Fatal("subscription was not created") + } + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("subagent subscription context was not canceled") + } +} + type fakeHandle struct { stream chan agentadaptor.StreamPayload events chan agentadaptor.RunEvent diff --git a/pkg/hosttools/a2adelegation/delegator.go b/pkg/hosttools/a2adelegation/delegator.go index 3775c31..d3059ef 100644 --- a/pkg/hosttools/a2adelegation/delegator.go +++ b/pkg/hosttools/a2adelegation/delegator.go @@ -11,10 +11,15 @@ import ( clienta2a "github.com/agent-dance/agent-adaptor/pkg/clients/a2a" ) +type A2AStream interface { + Recv() (clienta2a.Event, error) + Close() error +} + type A2AClient interface { AgentCard(ctx context.Context) (clienta2a.AgentCard, error) Send(ctx context.Context, req clienta2a.SendRequest) (clienta2a.Task, error) - SendStream(ctx context.Context, req clienta2a.SendRequest) (*clienta2a.Stream, error) + SendStream(ctx context.Context, req clienta2a.SendRequest) (A2AStream, error) GetTask(ctx context.Context, req clienta2a.GetTaskRequest) (clienta2a.Task, error) CancelTask(ctx context.Context, req clienta2a.CancelTaskRequest) (clienta2a.Task, error) } @@ -55,6 +60,13 @@ func (d *Delegator) Delegate(ctx context.Context, req DelegationRequest) (Delega } baseResult := DelegationResult{DelegationID: delegationID, Agent: spec.Key, RemoteProtocol: ProtocolA2A, Status: "running"} + timeout := clampTimeout(req.Timeout, spec.Policy.MaxTimeout) + if timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + client := d.clientFor(spec) card, err := client.AgentCard(ctx) if err != nil && spec.AgentCard == nil { @@ -75,13 +87,6 @@ func (d *Delegator) Delegate(ctx context.Context, req DelegationRequest) (Delega return baseResult, derr } - timeout := clampTimeout(req.Timeout, spec.Policy.MaxTimeout) - if timeout > 0 { - var cancel context.CancelFunc - ctx, cancel = context.WithTimeout(ctx, timeout) - defer cancel() - } - message, err := messageForDelegation(req) if err != nil { derr := delegationErr(err) @@ -120,20 +125,58 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe defer stream.Close() mapper := newEventMapper(baseEvent) var lastTask clienta2a.Task + lastTaskID := "" + cancelKnownTask := func(publish bool) { + if lastTaskID == "" { + return + } + if publish { + d.cancelRemote(context.Background(), client, lastTaskID, send.Tenant, baseEvent) + return + } + d.cancelRemoteTask(context.Background(), client, lastTaskID, send.Tenant, baseEvent) + } for { - event, err := stream.Recv() + select { + case <-ctx.Done(): + cancelKnownTask(true) + derr := &DelegationError{Code: "cancelled", Message: ctx.Err().Error(), Retryable: true} + baseResult.Status = "cancelled" + baseResult.Error = derr + return baseResult, derr + default: + } + item := make(chan streamRecv, 1) + go func() { + event, err := stream.Recv() + item <- streamRecv{event: event, err: err} + }() + var event clienta2a.Event + select { + case <-ctx.Done(): + cancelKnownTask(true) + _ = stream.Close() + derr := &DelegationError{Code: "cancelled", Message: ctx.Err().Error(), Retryable: true} + baseResult.Status = "cancelled" + baseResult.Error = derr + return baseResult, derr + case received := <-item: + event = received.event + err = received.err + } if err == io.EOF { break } if err != nil { - if lastTask.ID != "" { - if recovered, ok := d.recoverTask(ctx, client, spec, lastTask.ID); ok { + if lastTaskID != "" { + if recovered, ok := d.recoverTask(ctx, client, lastTaskID, send.Tenant); ok { lastTask = recovered for _, ev := range mapper.taskEvents(recovered) { d.publish(ev) } return d.finishTask(baseEvent, baseResult, recovered, spec.Policy, maxArtifacts) } + cancelKnownTask(false) } derr := &DelegationError{Code: "stream_interrupted", Message: err.Error(), Retryable: true} d.publish(failedEvent(baseEvent, derr)) @@ -141,17 +184,36 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe baseResult.Error = derr return baseResult, derr } + if event.TaskID != "" { + lastTaskID = event.TaskID + } + if event.Message != nil && event.Message.TaskID != "" { + lastTaskID = event.Message.TaskID + } if event.Task != nil { lastTask = *event.Task + lastTaskID = event.Task.ID } for _, ev := range mapper.Map(event) { d.publish(ev) } + if event.Message != nil && event.Kind == clienta2a.EventTerminal { + baseResult.RemoteTaskID = event.TaskID + baseResult.RemoteContextID = event.ContextID + baseResult.Status = "completed" + baseResult.Summary = textFromMessage(*event.Message) + if baseResult.Summary != "" { + baseResult.Messages = append(baseResult.Messages, DelegationMessage{Role: event.Message.Role, Text: baseResult.Summary}) + } + terminal := mapper.terminalForState(event.TaskID, event.ContextID, clienta2a.TaskStateCompleted, event.Raw) + d.publish(terminal) + return baseResult, nil + } if event.Task != nil && executionFinalState(event.Task.Status.State) { return d.finishTask(baseEvent, baseResult, *event.Task, spec.Policy, maxArtifacts) } if event.Status != nil && executionFinalState(event.Status.State) { - task, ok := d.recoverTask(ctx, client, spec, event.TaskID) + task, ok := d.recoverTask(ctx, client, event.TaskID, send.Tenant) if ok { return d.finishTask(baseEvent, baseResult, task, spec.Policy, maxArtifacts) } @@ -164,6 +226,8 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe baseResult.Error = derr return baseResult, derr } + terminal := mapper.terminalForState(event.TaskID, event.ContextID, event.Status.State, event.Raw) + d.publish(terminal) if baseResult.Status != "completed" { baseResult.Error = &DelegationError{Code: errorCodeFromState(event.Status.State), Message: "remote task did not complete successfully", RemoteStatus: string(event.Status.State)} } @@ -173,6 +237,7 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe if lastTask.ID != "" && executionFinalState(lastTask.Status.State) { return d.finishTask(baseEvent, baseResult, lastTask, spec.Policy, maxArtifacts) } + cancelKnownTask(false) derr := &DelegationError{Code: "stream_interrupted", Message: "remote stream ended before terminal state", Retryable: true} d.publish(failedEvent(baseEvent, derr)) baseResult.Status = "failed" @@ -180,6 +245,11 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe return baseResult, derr } +type streamRecv struct { + event clienta2a.Event + err error +} + func (d *Delegator) delegatePolling(ctx context.Context, client A2AClient, spec RemoteAgentSpec, send clienta2a.SendRequest, baseEvent DelegationEvent, baseResult DelegationResult, maxArtifacts *int) (DelegationResult, error) { send.ReturnImmediately = true task, err := client.Send(ctx, send) @@ -236,18 +306,16 @@ func (d *Delegator) delegatePolling(ctx context.Context, client A2AClient, spec return baseResult, derr } -func (d *Delegator) recoverTask(ctx context.Context, client A2AClient, spec RemoteAgentSpec, taskID string) (clienta2a.Task, bool) { +func (d *Delegator) recoverTask(ctx context.Context, client A2AClient, taskID, tenant string) (clienta2a.Task, bool) { if taskID == "" { return clienta2a.Task{}, false } - task, err := client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: taskID, Tenant: spec.Tenant}) + task, err := client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: taskID, Tenant: tenant}) return task, err == nil && executionFinalState(task.Status.State) } func (d *Delegator) cancelRemote(ctx context.Context, client A2AClient, taskID, tenant string, base DelegationEvent) { - if taskID != "" { - _, _ = client.CancelTask(ctx, clienta2a.CancelTaskRequest{TaskID: taskID, Tenant: tenant, Metadata: map[string]any{"reason": "parent_cancelled", "delegation_id": base.DelegationID}}) - } + d.cancelRemoteTask(ctx, client, taskID, tenant, base) ev := base ev.Kind = DelegationCancelled ev.RemoteTaskID = taskID @@ -255,6 +323,12 @@ func (d *Delegator) cancelRemote(ctx context.Context, client A2AClient, taskID, d.publish(ev) } +func (d *Delegator) cancelRemoteTask(ctx context.Context, client A2AClient, taskID, tenant string, base DelegationEvent) { + if taskID != "" { + _, _ = client.CancelTask(ctx, clienta2a.CancelTaskRequest{TaskID: taskID, Tenant: tenant, Metadata: map[string]any{"reason": "parent_cancelled", "delegation_id": base.DelegationID}}) + } +} + func (d *Delegator) publish(ev DelegationEvent) { if d.Bus != nil { d.Bus.Publish(ev) @@ -265,14 +339,22 @@ func (d *Delegator) clientFor(spec RemoteAgentSpec) A2AClient { if d.NewClient != nil { return d.NewClient(spec) } - return clienta2a.New(clienta2a.Options{ + return a2aClientAdapter{Client: clienta2a.New(clienta2a.Options{ AgentCardURL: spec.AgentCardURL, Auth: spec.Auth, HTTPClient: spec.HTTPClient, TrustedAuthOrigins: spec.TrustedAuthOrigins, AcceptedOutputModes: spec.AcceptedOutputModes, PreferredTransports: spec.PreferredTransports, - }) + })} +} + +type a2aClientAdapter struct { + *clienta2a.Client +} + +func (c a2aClientAdapter) SendStream(ctx context.Context, req clienta2a.SendRequest) (A2AStream, error) { + return c.Client.SendStream(ctx, req) } func (d *Delegator) newID() string { @@ -352,6 +434,8 @@ func (d *Delegator) finishTask(baseEvent DelegationEvent, baseResult DelegationR result.Error = nil } result.Artifacts = limitArtifacts(result.Artifacts, maxArtifacts) + mapper := newEventMapper(baseEvent) + d.publish(mapper.terminalForState(task.ID, task.ContextID, task.Status.State, task.Raw)) return result, terminalError(task.Status.State, policy) } diff --git a/pkg/hosttools/a2adelegation/delegator_test.go b/pkg/hosttools/a2adelegation/delegator_test.go index 472fb34..8aa31a1 100644 --- a/pkg/hosttools/a2adelegation/delegator_test.go +++ b/pkg/hosttools/a2adelegation/delegator_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "io" "net/http" "net/http/httptest" "strings" @@ -105,13 +106,8 @@ func TestEventMapperMapsA2ABridgeArtifactsAndTerminal(t *testing.T) { t.Fatalf("unexpected mapped events: %#v", events) } - terminal := mapper.Map(clienta2a.Event{ - Kind: clienta2a.EventTerminal, - TaskID: "task-1", - ContextID: "ctx-1", - Status: &clienta2a.TaskStatus{State: clienta2a.TaskStateInputRequired}, - }) - if len(terminal) != 1 || terminal[0].Kind != DelegationInputRequired { + terminal := mapper.terminalForState("task-1", "ctx-1", clienta2a.TaskStateInputRequired, nil) + if terminal.Kind != DelegationInputRequired { t.Fatalf("input-required terminal mapping: %#v", terminal) } } @@ -238,23 +234,24 @@ func TestDelegatorInputRequiredPolicy(t *testing.T) { t.Parallel() card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} for _, tc := range []struct { - name string - allow bool - wantErr bool - wantStatus string + name string + allow bool + wantErr bool + wantStatus string + wantTerminal DelegationEventKind }{ - {name: "default rejects", allow: false, wantErr: true, wantStatus: "failed"}, - {name: "policy allows", allow: true, wantErr: false, wantStatus: "input_required"}, + {name: "default rejects", allow: false, wantErr: true, wantStatus: "failed", wantTerminal: DelegationFailed}, + {name: "policy allows", allow: true, wantErr: false, wantStatus: "input_required", wantTerminal: DelegationInputRequired}, } { - tc := tc t.Run(tc.name, func(t *testing.T) { t.Parallel() registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Policy: DelegationPolicy{AllowInputRequired: tc.allow}}) if err != nil { t.Fatalf("registry: %v", err) } + bus := NewEventBus(16) client := &fakeA2AClient{card: card, sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateInputRequired}}} - delegator := NewDelegator(registry, NewEventBus(16)) + delegator := NewDelegator(registry, bus) delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } delegator.NewID = func() string { return "del-1" } @@ -271,6 +268,10 @@ func TestDelegatorInputRequiredPolicy(t *testing.T) { if !tc.wantErr && result.Error != nil { t.Fatalf("allowed input-required should not set error: %#v", result.Error) } + replayed := drainBus(t, bus, "run-1", 3) + if replayed[len(replayed)-1].Kind != tc.wantTerminal { + t.Fatalf("terminal: got %#v want %s in %#v", replayed[len(replayed)-1], tc.wantTerminal, replayed) + } }) } } @@ -282,8 +283,9 @@ func TestDelegatorMaxArtifactBytesPolicy(t *testing.T) { if err != nil { t.Fatalf("registry: %v", err) } + bus := NewEventBus(16) client := &fakeA2AClient{card: card, sendTask: clienta2a.Task{ID: "task-1", ContextID: "ctx-1", Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}, Artifacts: []clienta2a.Artifact{{ID: "large", Name: "large.txt", Parts: []clienta2a.Part{{Kind: clienta2a.PartText, Text: "too large"}}}}}} - delegator := NewDelegator(registry, NewEventBus(16)) + delegator := NewDelegator(registry, bus) delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } delegator.NewID = func() string { return "del-1" } @@ -298,6 +300,10 @@ func TestDelegatorMaxArtifactBytesPolicy(t *testing.T) { if result.Status != "failed" { t.Fatalf("expected failed result, got %#v", result) } + replayed := drainBus(t, bus, "run-1", 4) + if replayed[len(replayed)-1].Kind != DelegationFailed { + t.Fatalf("policy failure should publish failed terminal last, got %#v", replayed) + } } func TestDelegatorSendsInputArtifactsAsURLParts(t *testing.T) { @@ -392,6 +398,132 @@ func TestDelegatorPollingCancellationCascadesToRemoteTask(t *testing.T) { } } +func TestDelegatorStreamingCancellationCascadesWithEffectiveTenant(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Tenant: "spec-tenant"}) + if err != nil { + t.Fatalf("registry: %v", err) + } + stream := &fakeA2AStream{events: make(chan streamRecv, 1), closed: make(chan struct{})} + stream.events <- streamRecv{event: clienta2a.Event{Kind: clienta2a.EventStatus, TaskID: "task-1", ContextID: "ctx-1", Status: &clienta2a.TaskStatus{State: clienta2a.TaskStateWorking}}} + client := &fakeA2AClient{card: card, stream: stream} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + result, err := delegator.Delegate(ctx, DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Stream: true, Tenant: "call-tenant"}) + if err == nil { + t.Fatal("expected streaming cancellation error") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "cancelled" { + t.Fatalf("expected cancelled error, got %T %[1]v", err) + } + if result.Status != "cancelled" { + t.Fatalf("expected cancelled result, got %#v", result) + } + if client.cancelCalls != 1 || client.lastCancel.TaskID != "task-1" || client.lastCancel.Tenant != "call-tenant" { + t.Fatalf("expected streaming remote cancel with effective tenant, calls=%d req=%#v", client.cancelCalls, client.lastCancel) + } +} + +func TestDelegatorStreamingRecoveryUsesEffectiveTenant(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card, Tenant: "spec-tenant"}) + if err != nil { + t.Fatalf("registry: %v", err) + } + stream := &fakeA2AStream{events: make(chan streamRecv, 1), closed: make(chan struct{})} + stream.events <- streamRecv{event: clienta2a.Event{Kind: clienta2a.EventStatus, TaskID: "task-1", ContextID: "ctx-1", Status: &clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}}} + close(stream.events) + client := &fakeA2AClient{ + card: card, + stream: stream, + getTasks: []clienta2a.Task{{ + ID: "task-1", + ContextID: "ctx-1", + Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCompleted}, + Messages: []clienta2a.Message{{Role: "agent", Parts: []clienta2a.Part{{Kind: clienta2a.PartText, Text: "done"}}}}, + }}, + } + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Stream: true, Tenant: "call-tenant"}) + if err != nil { + t.Fatalf("delegate: %v", err) + } + if result.Status != "completed" || result.Summary != "done" { + t.Fatalf("unexpected result: %#v", result) + } + if client.getCalls != 1 || client.lastGet.TaskID != "task-1" || client.lastGet.Tenant != "call-tenant" { + t.Fatalf("expected recovery get with effective tenant, calls=%d req=%#v", client.getCalls, client.lastGet) + } +} + +func TestDelegatorStreamingTerminalMessageCompletes(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + stream := &fakeA2AStream{events: make(chan streamRecv, 1), closed: make(chan struct{})} + stream.events <- streamRecv{event: clienta2a.Event{ + Kind: clienta2a.EventTerminal, + TaskID: "task-1", + ContextID: "ctx-1", + Message: &clienta2a.Message{Role: "agent", TaskID: "task-1", ContextID: "ctx-1", Parts: []clienta2a.Part{{Kind: clienta2a.PartText, Text: "done"}}}, + }} + close(stream.events) + bus := NewEventBus(16) + client := &fakeA2AClient{card: card, stream: stream} + delegator := NewDelegator(registry, bus) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Stream: true}) + if err != nil { + t.Fatalf("delegate: %v", err) + } + if result.Status != "completed" || result.Summary != "done" || client.cancelCalls != 0 { + t.Fatalf("unexpected streaming terminal message result=%#v cancelCalls=%d", result, client.cancelCalls) + } + replayed := drainBus(t, bus, "run-1", 4) + if replayed[len(replayed)-1].Kind != DelegationFinished { + t.Fatalf("expected finished terminal, got %#v", replayed) + } +} + +func TestDelegatorTimeoutCoversAgentCardDiscovery(t *testing.T) { + t.Parallel() + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCardURL: "https://agent.example/card"}) + if err != nil { + t.Fatalf("registry: %v", err) + } + client := &fakeA2AClient{cardErr: context.DeadlineExceeded} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Timeout: time.Millisecond}) + if err == nil { + t.Fatal("expected agent card timeout") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "agent_unavailable" { + t.Fatalf("expected agent_unavailable timeout, got %T %[1]v", err) + } + if result.Status != "failed" { + t.Fatalf("expected failed result, got %#v", result) + } +} + func TestEventBusDeduplicatesTerminalEventsAndReplays(t *testing.T) { t.Parallel() bus := NewEventBus(8) @@ -450,6 +582,35 @@ func TestEventBusPublishDoesNotBlockOnFullSubscriber(t *testing.T) { } } +func TestEventBusClearRunDropsStateAndClosesSubscribers(t *testing.T) { + t.Parallel() + bus := NewEventBus(8) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := bus.SubscribeRun(ctx, "run-1") + bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationStarted}) + bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationFinished}) + bus.ClearRun("run-1") + deadline := time.After(time.Second) + for { + select { + case _, ok := <-ch: + if !ok { + goto closed + } + case <-deadline: + t.Fatal("subscriber did not close after ClearRun") + } + } +closed: + if replayed := drainAvailableBus(t, bus, "run-1"); len(replayed) != 0 { + t.Fatalf("expected replay to be cleared, got %#v", replayed) + } + if !bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationFailed}) { + t.Fatal("terminal state should be cleared for run") + } +} + func drainBus(t *testing.T, bus *EventBus, runID string, want int) []DelegationEvent { t.Helper() ctx, cancel := context.WithCancel(context.Background()) @@ -467,6 +628,22 @@ func drainBus(t *testing.T, bus *EventBus, runID string, want int) []DelegationE return out } +func drainAvailableBus(t *testing.T, bus *EventBus, runID string) []DelegationEvent { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := bus.SubscribeRun(ctx, runID) + out := []DelegationEvent{} + for { + select { + case ev := <-ch: + out = append(out, ev) + default: + return out + } + } +} + func TestMCPServerDelegatesToolCallAndReturnsStructuredResult(t *testing.T) { t.Parallel() card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} @@ -536,22 +713,32 @@ func TestMCPServerRejectsMissingBearerToken(t *testing.T) { } type fakeA2AClient struct { - card clienta2a.AgentCard + card clienta2a.AgentCard + cardErr error sendTask clienta2a.Task lastSend clienta2a.SendRequest sendCalls int + stream A2AStream streamCalls int getTasks []clienta2a.Task getErr error + lastGet clienta2a.GetTaskRequest + getCalls int cancelCalls int lastCancel clienta2a.CancelTaskRequest } -func (f *fakeA2AClient) AgentCard(context.Context) (clienta2a.AgentCard, error) { return f.card, nil } +func (f *fakeA2AClient) AgentCard(ctx context.Context) (clienta2a.AgentCard, error) { + if f.cardErr != nil { + <-ctx.Done() + return clienta2a.AgentCard{}, ctx.Err() + } + return f.card, nil +} func (f *fakeA2AClient) Send(_ context.Context, req clienta2a.SendRequest) (clienta2a.Task, error) { f.sendCalls++ @@ -559,12 +746,17 @@ func (f *fakeA2AClient) Send(_ context.Context, req clienta2a.SendRequest) (clie return f.sendTask, nil } -func (f *fakeA2AClient) SendStream(context.Context, clienta2a.SendRequest) (*clienta2a.Stream, error) { +func (f *fakeA2AClient) SendStream(context.Context, clienta2a.SendRequest) (A2AStream, error) { f.streamCalls++ + if f.stream != nil { + return f.stream, nil + } return nil, errors.New("stream unavailable") } -func (f *fakeA2AClient) GetTask(context.Context, clienta2a.GetTaskRequest) (clienta2a.Task, error) { +func (f *fakeA2AClient) GetTask(_ context.Context, req clienta2a.GetTaskRequest) (clienta2a.Task, error) { + f.getCalls++ + f.lastGet = req if f.getErr != nil { return clienta2a.Task{}, f.getErr } @@ -581,3 +773,25 @@ func (f *fakeA2AClient) CancelTask(_ context.Context, req clienta2a.CancelTaskRe f.lastCancel = req return clienta2a.Task{ID: req.TaskID, Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCanceled}}, nil } + +type fakeA2AStream struct { + events chan streamRecv + closed chan struct{} +} + +func (s *fakeA2AStream) Recv() (clienta2a.Event, error) { + item, ok := <-s.events + if !ok { + return clienta2a.Event{}, io.EOF + } + return item.event, item.err +} + +func (s *fakeA2AStream) Close() error { + select { + case <-s.closed: + default: + close(s.closed) + } + return nil +} diff --git a/pkg/hosttools/a2adelegation/events.go b/pkg/hosttools/a2adelegation/events.go index 99935ed..01699c7 100644 --- a/pkg/hosttools/a2adelegation/events.go +++ b/pkg/hosttools/a2adelegation/events.go @@ -64,6 +64,20 @@ func (b *EventBus) Publish(ev DelegationEvent) bool { return true } +func (b *EventBus) ClearRun(runID string) { + if b == nil || runID == "" { + return + } + b.mu.Lock() + defer b.mu.Unlock() + for ch := range b.subscribers[runID] { + closeDelegationChannel(ch) + } + delete(b.subscribers, runID) + delete(b.replay, runID) + delete(b.terminal, runID) +} + func (b *EventBus) SubscribeRun(ctx context.Context, runID string) <-chan DelegationEvent { if b == nil || runID == "" { out := make(chan DelegationEvent) @@ -90,11 +104,16 @@ func (b *EventBus) SubscribeRun(ctx context.Context, runID string) <-chan Delega delete(b.subscribers, runID) } b.mu.Unlock() - close(out) + closeDelegationChannel(out) }() return out } +func closeDelegationChannel(ch chan DelegationEvent) { + defer func() { _ = recover() }() + close(ch) +} + func isTerminal(kind DelegationEventKind) bool { switch kind { case DelegationFinished, DelegationFailed, DelegationCancelled, DelegationInputRequired: diff --git a/pkg/hosttools/a2adelegation/mapping.go b/pkg/hosttools/a2adelegation/mapping.go index 9b5fd2b..5a65410 100644 --- a/pkg/hosttools/a2adelegation/mapping.go +++ b/pkg/hosttools/a2adelegation/mapping.go @@ -50,7 +50,9 @@ func (m *eventMapper) Map(event clienta2a.Event) []DelegationEvent { out = append(out, m.artifactEvent(event.TaskID, event.ContextID, *event.Artifact)) } case clienta2a.EventTerminal: - out = append(out, m.terminalEvents(event)...) + if event.Message != nil { + out = append(out, m.messageEvents(*event.Message)...) + } } return out } @@ -63,9 +65,6 @@ func (m *eventMapper) taskEvents(task clienta2a.Task) []DelegationEvent { for _, artifact := range task.Artifacts { out = append(out, m.artifactEvent(task.ID, task.ContextID, artifact)) } - if executionFinalState(task.Status.State) { - out = append(out, m.terminalForState(task.ID, task.ContextID, task.Status.State, task.Raw)) - } return out } From c23cdc041803520631cf87dc6dfffac3022e5842 Mon Sep 17 00:00:00 2001 From: buthim Date: Tue, 7 Jul 2026 15:45:22 +0800 Subject: [PATCH 3/6] fix: stabilize streaming delegation cancellation Co-Authored-By: Claude Opus 4.8 --- pkg/hosttools/a2adelegation/delegator.go | 46 +++++++------ pkg/hosttools/a2adelegation/delegator_test.go | 67 ++++++++++++++++++- 2 files changed, 92 insertions(+), 21 deletions(-) diff --git a/pkg/hosttools/a2adelegation/delegator.go b/pkg/hosttools/a2adelegation/delegator.go index d3059ef..60c37b8 100644 --- a/pkg/hosttools/a2adelegation/delegator.go +++ b/pkg/hosttools/a2adelegation/delegator.go @@ -3,6 +3,7 @@ package a2adelegation import ( "context" "encoding/json" + "errors" "fmt" "io" "strings" @@ -126,24 +127,26 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe mapper := newEventMapper(baseEvent) var lastTask clienta2a.Task lastTaskID := "" - cancelKnownTask := func(publish bool) { - if lastTaskID == "" { - return - } - if publish { + cancelResult := func() (DelegationResult, error) { + if lastTaskID != "" { d.cancelRemote(context.Background(), client, lastTaskID, send.Tenant, baseEvent) - return + } else { + d.publish(cancelledEvent(baseEvent, "")) + } + derr := &DelegationError{Code: "cancelled", Message: ctx.Err().Error(), Retryable: true} + baseResult.Status = "cancelled" + baseResult.Error = derr + return baseResult, derr + } + cancelKnownTask := func() { + if lastTaskID != "" { + d.cancelRemoteTask(context.Background(), client, lastTaskID, send.Tenant, baseEvent) } - d.cancelRemoteTask(context.Background(), client, lastTaskID, send.Tenant, baseEvent) } for { select { case <-ctx.Done(): - cancelKnownTask(true) - derr := &DelegationError{Code: "cancelled", Message: ctx.Err().Error(), Retryable: true} - baseResult.Status = "cancelled" - baseResult.Error = derr - return baseResult, derr + return cancelResult() default: } item := make(chan streamRecv, 1) @@ -154,12 +157,8 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe var event clienta2a.Event select { case <-ctx.Done(): - cancelKnownTask(true) _ = stream.Close() - derr := &DelegationError{Code: "cancelled", Message: ctx.Err().Error(), Retryable: true} - baseResult.Status = "cancelled" - baseResult.Error = derr - return baseResult, derr + return cancelResult() case received := <-item: event = received.event err = received.err @@ -168,6 +167,9 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe break } if err != nil { + if ctx.Err() != nil && (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) { + return cancelResult() + } if lastTaskID != "" { if recovered, ok := d.recoverTask(ctx, client, lastTaskID, send.Tenant); ok { lastTask = recovered @@ -176,7 +178,7 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe } return d.finishTask(baseEvent, baseResult, recovered, spec.Policy, maxArtifacts) } - cancelKnownTask(false) + cancelKnownTask() } derr := &DelegationError{Code: "stream_interrupted", Message: err.Error(), Retryable: true} d.publish(failedEvent(baseEvent, derr)) @@ -237,7 +239,7 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe if lastTask.ID != "" && executionFinalState(lastTask.Status.State) { return d.finishTask(baseEvent, baseResult, lastTask, spec.Policy, maxArtifacts) } - cancelKnownTask(false) + cancelKnownTask() derr := &DelegationError{Code: "stream_interrupted", Message: "remote stream ended before terminal state", Retryable: true} d.publish(failedEvent(baseEvent, derr)) baseResult.Status = "failed" @@ -316,11 +318,15 @@ func (d *Delegator) recoverTask(ctx context.Context, client A2AClient, taskID, t func (d *Delegator) cancelRemote(ctx context.Context, client A2AClient, taskID, tenant string, base DelegationEvent) { d.cancelRemoteTask(ctx, client, taskID, tenant, base) + d.publish(cancelledEvent(base, taskID)) +} + +func cancelledEvent(base DelegationEvent, taskID string) DelegationEvent { ev := base ev.Kind = DelegationCancelled ev.RemoteTaskID = taskID ev.Status = "cancelled" - d.publish(ev) + return ev } func (d *Delegator) cancelRemoteTask(ctx context.Context, client A2AClient, taskID, tenant string, base DelegationEvent) { diff --git a/pkg/hosttools/a2adelegation/delegator_test.go b/pkg/hosttools/a2adelegation/delegator_test.go index 8aa31a1..54ca060 100644 --- a/pkg/hosttools/a2adelegation/delegator_test.go +++ b/pkg/hosttools/a2adelegation/delegator_test.go @@ -407,8 +407,9 @@ func TestDelegatorStreamingCancellationCascadesWithEffectiveTenant(t *testing.T) } stream := &fakeA2AStream{events: make(chan streamRecv, 1), closed: make(chan struct{})} stream.events <- streamRecv{event: clienta2a.Event{Kind: clienta2a.EventStatus, TaskID: "task-1", ContextID: "ctx-1", Status: &clienta2a.TaskStatus{State: clienta2a.TaskStateWorking}}} + bus := NewEventBus(16) client := &fakeA2AClient{card: card, stream: stream} - delegator := NewDelegator(registry, NewEventBus(16)) + delegator := NewDelegator(registry, bus) delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } delegator.NewID = func() string { return "del-1" } ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) @@ -428,6 +429,10 @@ func TestDelegatorStreamingCancellationCascadesWithEffectiveTenant(t *testing.T) if client.cancelCalls != 1 || client.lastCancel.TaskID != "task-1" || client.lastCancel.Tenant != "call-tenant" { t.Fatalf("expected streaming remote cancel with effective tenant, calls=%d req=%#v", client.cancelCalls, client.lastCancel) } + replayed := drainBus(t, bus, "run-1", 3) + if replayed[len(replayed)-1].Kind != DelegationCancelled { + t.Fatalf("expected cancelled terminal, got %#v", replayed) + } } func TestDelegatorStreamingRecoveryUsesEffectiveTenant(t *testing.T) { @@ -466,6 +471,66 @@ func TestDelegatorStreamingRecoveryUsesEffectiveTenant(t *testing.T) { } } +func TestDelegatorStreamingCancelBeforeTaskIDPublishesTerminal(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + stream := &fakeA2AStream{events: make(chan streamRecv), closed: make(chan struct{})} + bus := NewEventBus(16) + client := &fakeA2AClient{card: card, stream: stream} + delegator := NewDelegator(registry, bus) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond) + defer cancel() + + result, err := delegator.Delegate(ctx, DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Stream: true}) + if err == nil { + t.Fatal("expected cancellation before first task id") + } + if result.Status != "cancelled" || client.cancelCalls != 0 { + t.Fatalf("unexpected result=%#v cancelCalls=%d", result, client.cancelCalls) + } + replayed := drainBus(t, bus, "run-1", 1) + if replayed[0].Kind != DelegationCancelled || replayed[0].RemoteTaskID != "" { + t.Fatalf("expected cancelled terminal without remote task id, got %#v", replayed) + } +} + +func TestDelegatorStreamingRecvContextErrorReportsCancelled(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + stream := &fakeA2AStream{events: make(chan streamRecv, 1), closed: make(chan struct{})} + stream.events <- streamRecv{err: context.DeadlineExceeded} + bus := NewEventBus(16) + client := &fakeA2AClient{card: card, stream: stream} + delegator := NewDelegator(registry, bus) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + result, err := delegator.Delegate(ctx, DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Stream: true}) + if err == nil { + t.Fatal("expected cancelled stream recv error") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "cancelled" || result.Status != "cancelled" { + t.Fatalf("expected cancelled result/error, result=%#v err=%T %[2]v", result, err) + } + replayed := drainBus(t, bus, "run-1", 1) + if replayed[0].Kind != DelegationCancelled { + t.Fatalf("expected cancelled terminal, got %#v", replayed) + } +} + func TestDelegatorStreamingTerminalMessageCompletes(t *testing.T) { t.Parallel() card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} From 6ddbcc57f719d34e4db5aae88461155068b65bc1 Mon Sep 17 00:00:00 2001 From: buthim Date: Tue, 7 Jul 2026 16:00:34 +0800 Subject: [PATCH 4/6] fix: harden delegation runtime edge cases Co-Authored-By: Claude Opus 4.8 --- pkg/bridges/sse/handler.go | 5 +- pkg/hosttools/a2adelegation/delegator.go | 28 +++++-- pkg/hosttools/a2adelegation/delegator_test.go | 77 ++++++++++++++++++- pkg/hosttools/a2adelegation/events.go | 24 +++++- 4 files changed, 119 insertions(+), 15 deletions(-) diff --git a/pkg/bridges/sse/handler.go b/pkg/bridges/sse/handler.go index b4d6a13..b9d2cdc 100644 --- a/pkg/bridges/sse/handler.go +++ b/pkg/bridges/sse/handler.go @@ -40,6 +40,7 @@ import ( "strings" "time" + aguievents "github.com/ag-ui-protocol/ag-ui/sdks/community/go/pkg/core/events" aguisse "github.com/ag-ui-protocol/ag-ui/sdks/community/go/pkg/encoding/sse" agentadaptor "github.com/agent-dance/agent-adaptor" @@ -245,9 +246,11 @@ func decodeRawRequest(r *http.Request) (*RawRequest, error) { func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, w io.Writer, handle agentadaptor.RunHandle, opts Options) error { switch opts.Protocol { case AGUI: - out := agui.WrapWithContext(ctx, handle) + var out <-chan aguievents.Event if opts.SubagentBus != nil { out = subagentstream.WrapAGUI(ctx, handle, subagentstream.MuxOptions{Bus: opts.SubagentBus}) + } else { + out = agui.WrapWithContext(ctx, handle) } for ev := range out { select { diff --git a/pkg/hosttools/a2adelegation/delegator.go b/pkg/hosttools/a2adelegation/delegator.go index 60c37b8..3d6fc81 100644 --- a/pkg/hosttools/a2adelegation/delegator.go +++ b/pkg/hosttools/a2adelegation/delegator.go @@ -2,16 +2,23 @@ package a2adelegation import ( "context" + "crypto/rand" + "encoding/hex" "encoding/json" "errors" "fmt" "io" "strings" + "sync/atomic" "time" clienta2a "github.com/agent-dance/agent-adaptor/pkg/clients/a2a" ) +const remoteCancelTimeout = 5 * time.Second + +var delegationIDCounter atomic.Uint64 + type A2AStream interface { Recv() (clienta2a.Event, error) Close() error @@ -129,7 +136,7 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe lastTaskID := "" cancelResult := func() (DelegationResult, error) { if lastTaskID != "" { - d.cancelRemote(context.Background(), client, lastTaskID, send.Tenant, baseEvent) + d.cancelRemote(ctx, client, lastTaskID, send.Tenant, baseEvent) } else { d.publish(cancelledEvent(baseEvent, "")) } @@ -140,7 +147,7 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe } cancelKnownTask := func() { if lastTaskID != "" { - d.cancelRemoteTask(context.Background(), client, lastTaskID, send.Tenant, baseEvent) + d.cancelRemoteTask(ctx, client, lastTaskID, send.Tenant, baseEvent) } } for { @@ -283,7 +290,7 @@ func (d *Delegator) delegatePolling(ctx context.Context, client A2AClient, spec for i := 0; i < maxPolls; i++ { select { case <-ctx.Done(): - d.cancelRemote(context.Background(), client, task.ID, send.Tenant, baseEvent) + d.cancelRemote(ctx, client, task.ID, send.Tenant, baseEvent) derr := &DelegationError{Code: "cancelled", Message: ctx.Err().Error(), Retryable: true} baseResult.Status = "cancelled" baseResult.Error = derr @@ -301,7 +308,7 @@ func (d *Delegator) delegatePolling(ctx context.Context, client A2AClient, spec return d.finishTask(baseEvent, baseResult, task, spec.Policy, maxArtifacts) } } - d.cancelRemote(context.Background(), client, task.ID, send.Tenant, baseEvent) + d.cancelRemote(ctx, client, task.ID, send.Tenant, baseEvent) derr := &DelegationError{Code: "remote_timeout", Message: "remote task did not finish before timeout", Retryable: true, RemoteStatus: string(task.Status.State)} baseResult.Status = "failed" baseResult.Error = derr @@ -330,9 +337,12 @@ func cancelledEvent(base DelegationEvent, taskID string) DelegationEvent { } func (d *Delegator) cancelRemoteTask(ctx context.Context, client A2AClient, taskID, tenant string, base DelegationEvent) { - if taskID != "" { - _, _ = client.CancelTask(ctx, clienta2a.CancelTaskRequest{TaskID: taskID, Tenant: tenant, Metadata: map[string]any{"reason": "parent_cancelled", "delegation_id": base.DelegationID}}) + if taskID == "" { + return } + cancelCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), remoteCancelTimeout) + defer cancel() + _, _ = client.CancelTask(cancelCtx, clienta2a.CancelTaskRequest{TaskID: taskID, Tenant: tenant, Metadata: map[string]any{"reason": "parent_cancelled", "delegation_id": base.DelegationID}}) } func (d *Delegator) publish(ev DelegationEvent) { @@ -367,7 +377,11 @@ func (d *Delegator) newID() string { if d != nil && d.NewID != nil { return d.NewID() } - return "del-" + time.Now().UTC().Format("20060102150405.000000000") + var random [8]byte + if _, err := rand.Read(random[:]); err == nil { + return "del-" + hex.EncodeToString(random[:]) + } + return fmt.Sprintf("del-%d", delegationIDCounter.Add(1)) } func messageForDelegation(req DelegationRequest) (clienta2a.Message, error) { diff --git a/pkg/hosttools/a2adelegation/delegator_test.go b/pkg/hosttools/a2adelegation/delegator_test.go index 54ca060..716af3a 100644 --- a/pkg/hosttools/a2adelegation/delegator_test.go +++ b/pkg/hosttools/a2adelegation/delegator_test.go @@ -435,6 +435,31 @@ func TestDelegatorStreamingCancellationCascadesWithEffectiveTenant(t *testing.T) } } +func TestDelegatorStreamingCancellationUsesBoundedRemoteCancelContext(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + stream := &fakeA2AStream{events: make(chan streamRecv, 1), closed: make(chan struct{})} + stream.events <- streamRecv{event: clienta2a.Event{Kind: clienta2a.EventStatus, TaskID: "task-1", ContextID: "ctx-1", Status: &clienta2a.TaskStatus{State: clienta2a.TaskStateWorking}}} + client := &fakeA2AClient{card: card, stream: stream} + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } + delegator.NewID = func() string { return "del-1" } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + _, err = delegator.Delegate(ctx, DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", Stream: true}) + if err == nil { + t.Fatal("expected streaming cancellation error") + } + if !client.cancelHadDeadline { + t.Fatalf("expected remote cancel to use bounded context, got calls=%d", client.cancelCalls) + } +} + func TestDelegatorStreamingRecoveryUsesEffectiveTenant(t *testing.T) { t.Parallel() card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: true}} @@ -647,6 +672,34 @@ func TestEventBusPublishDoesNotBlockOnFullSubscriber(t *testing.T) { } } +func TestEventBusTerminalDeliveryDropsOldestWhenSubscriberFull(t *testing.T) { + t.Parallel() + bus := NewEventBus(0) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := bus.SubscribeRun(ctx, "run-1") + for i := 0; i < subscriberBuffer; i++ { + bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationStatus, Status: "working"}) + } + if !bus.Publish(DelegationEvent{RunID: "run-1", DelegationID: "del-1", Kind: DelegationFinished, Status: "completed"}) { + t.Fatal("terminal publish should be accepted") + } + seenTerminal := false + for i := 0; i < subscriberBuffer; i++ { + select { + case ev := <-ch: + if ev.Kind == DelegationFinished { + seenTerminal = true + } + case <-time.After(time.Second): + t.Fatalf("timeout draining subscriber %d/%d", i, subscriberBuffer) + } + } + if !seenTerminal { + t.Fatal("terminal event should be enqueued by dropping an older event") + } +} + func TestEventBusClearRunDropsStateAndClosesSubscribers(t *testing.T) { t.Parallel() bus := NewEventBus(8) @@ -709,6 +762,22 @@ func drainAvailableBus(t *testing.T, bus *EventBus, runID string) []DelegationEv } } +func TestDelegatorDefaultIDsAreUnique(t *testing.T) { + t.Parallel() + delegator := NewDelegator(nil, nil) + seen := map[string]struct{}{} + for i := 0; i < 256; i++ { + id := delegator.newID() + if id == "" { + t.Fatal("newID returned empty string") + } + if _, ok := seen[id]; ok { + t.Fatalf("duplicate delegation id %q", id) + } + seen[id] = struct{}{} + } +} + func TestMCPServerDelegatesToolCallAndReturnsStructuredResult(t *testing.T) { t.Parallel() card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} @@ -793,8 +862,9 @@ type fakeA2AClient struct { lastGet clienta2a.GetTaskRequest getCalls int - cancelCalls int - lastCancel clienta2a.CancelTaskRequest + cancelCalls int + lastCancel clienta2a.CancelTaskRequest + cancelHadDeadline bool } func (f *fakeA2AClient) AgentCard(ctx context.Context) (clienta2a.AgentCard, error) { @@ -833,9 +903,10 @@ func (f *fakeA2AClient) GetTask(_ context.Context, req clienta2a.GetTaskRequest) return task, nil } -func (f *fakeA2AClient) CancelTask(_ context.Context, req clienta2a.CancelTaskRequest) (clienta2a.Task, error) { +func (f *fakeA2AClient) CancelTask(ctx context.Context, req clienta2a.CancelTaskRequest) (clienta2a.Task, error) { f.cancelCalls++ f.lastCancel = req + _, f.cancelHadDeadline = ctx.Deadline() return clienta2a.Task{ID: req.TaskID, Status: clienta2a.TaskStatus{State: clienta2a.TaskStateCanceled}}, nil } diff --git a/pkg/hosttools/a2adelegation/events.go b/pkg/hosttools/a2adelegation/events.go index 01699c7..9dc9866 100644 --- a/pkg/hosttools/a2adelegation/events.go +++ b/pkg/hosttools/a2adelegation/events.go @@ -56,14 +56,30 @@ func (b *EventBus) Publish(ev DelegationEvent) bool { b.replay[ev.RunID] = buf } for ch := range b.subscribers[ev.RunID] { - select { - case ch <- ev: - default: - } + deliverSubscriber(ch, ev) } return true } +func deliverSubscriber(ch chan DelegationEvent, ev DelegationEvent) { + select { + case ch <- ev: + return + default: + } + if !isTerminal(ev.Kind) { + return + } + select { + case <-ch: + default: + } + select { + case ch <- ev: + default: + } +} + func (b *EventBus) ClearRun(runID string) { if b == nil || runID == "" { return From 52dd4e6e443f47ccb9cfa674631fbb6c1b799c5c Mon Sep 17 00:00:00 2001 From: buthim Date: Tue, 7 Jul 2026 16:32:51 +0800 Subject: [PATCH 5/6] fix: address delegation review hardening Co-Authored-By: Claude Opus 4.8 --- pkg/bridges/sse/handler.go | 43 +++++++++++++------ pkg/bridges/subagentstream/mux.go | 3 -- pkg/bridges/subagentstream/mux_test.go | 19 ++++++++ pkg/hosttools/a2adelegation/delegator.go | 10 +++++ pkg/hosttools/a2adelegation/delegator_test.go | 23 ++++++++++ 5 files changed, 82 insertions(+), 16 deletions(-) diff --git a/pkg/bridges/sse/handler.go b/pkg/bridges/sse/handler.go index b9d2cdc..91b994d 100644 --- a/pkg/bridges/sse/handler.go +++ b/pkg/bridges/sse/handler.go @@ -38,6 +38,7 @@ import ( "io" "net/http" "strings" + "sync" "time" aguievents "github.com/ag-ui-protocol/ag-ui/sdks/community/go/pkg/core/events" @@ -158,20 +159,25 @@ func Handler(sdk agentadaptor.SDK, opts Options) http.Handler { w.WriteHeader(http.StatusOK) flusher, _ := w.(http.Flusher) + writeMu := &sync.Mutex{} if flusher != nil { + writeMu.Lock() flusher.Flush() + writeMu.Unlock() } pingCtx, cancelPing := context.WithCancel(r.Context()) defer cancelPing() if opts.KeepAlivePing > 0 && flusher != nil { - go runKeepAlive(pingCtx, w, flusher, opts.KeepAlivePing) + go runKeepAlive(pingCtx, writeMu, w, flusher, opts.KeepAlivePing) } ctx := r.Context() - err = streamEvents(ctx, writer, w, handle, opts) + err = streamEvents(ctx, writer, writeMu, w, handle, opts) if err != nil && !errors.Is(err, context.Canceled) { + writeMu.Lock() _ = writer.WriteErrorEvent(ctx, w, err, "") + writeMu.Unlock() } }) } @@ -243,7 +249,7 @@ func decodeRawRequest(r *http.Request) (*RawRequest, error) { // streamEvents drains the handle's streams into SSE frames according to // the chosen protocol. It returns when the stream is exhausted or the // context is cancelled. -func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, w io.Writer, handle agentadaptor.RunHandle, opts Options) error { +func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, writeMu *sync.Mutex, w io.Writer, handle agentadaptor.RunHandle, opts Options) error { switch opts.Protocol { case AGUI: var out <-chan aguievents.Event @@ -258,13 +264,16 @@ func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, w io.Writer, h return ctx.Err() default: } - if err := writer.WriteEvent(ctx, w, ev); err != nil { + writeMu.Lock() + err := writer.WriteEvent(ctx, w, ev) + writeMu.Unlock() + if err != nil { return err } } return nil case Raw: - return streamRaw(ctx, w, handle) + return streamRaw(ctx, writeMu, w, handle) default: return fmt.Errorf("sse: unknown protocol %d", opts.Protocol) } @@ -274,7 +283,7 @@ func streamEvents(ctx context.Context, writer *aguisse.SSEWriter, w io.Writer, h // HITL events are renamed to decision.request / decision.resolved and their // body is the corresponding structured payload (HITLRequestedPayload / // HITLResolvedPayload) — see docs/workstream-hitl-v2.md §6.2. -func streamRaw(ctx context.Context, w io.Writer, handle agentadaptor.RunHandle) error { +func streamRaw(ctx context.Context, writeMu *sync.Mutex, w io.Writer, handle agentadaptor.RunHandle) error { flusher, _ := w.(http.Flusher) done := make(chan agentadaptor.RunResult, 1) go func() { @@ -304,12 +313,15 @@ func streamRaw(ctx context.Context, w io.Writer, handle agentadaptor.RunHandle) if id == 0 { id = p.Sequence } - if _, err := fmt.Fprintf(w, "event: %s\nid: %d\ndata: %s\n\n", escapeEventName(name), id, payload); err != nil { - return fmt.Errorf("sse: write: %w", err) - } - if flusher != nil { + writeMu.Lock() + _, err = fmt.Fprintf(w, "event: %s\nid: %d\ndata: %s\n\n", escapeEventName(name), id, payload) + if err == nil && flusher != nil { flusher.Flush() } + writeMu.Unlock() + if err != nil { + return fmt.Errorf("sse: write: %w", err) + } } } } @@ -332,7 +344,7 @@ func rawFrameFor(p agentadaptor.StreamPayload) (string, any) { return string(p.Kind), p } -func runKeepAlive(ctx context.Context, w io.Writer, flusher http.Flusher, interval time.Duration) { +func runKeepAlive(ctx context.Context, writeMu *sync.Mutex, w io.Writer, flusher http.Flusher, interval time.Duration) { ticker := time.NewTicker(interval) defer ticker.Stop() for { @@ -340,10 +352,15 @@ func runKeepAlive(ctx context.Context, w io.Writer, flusher http.Flusher, interv case <-ctx.Done(): return case <-ticker.C: - if _, err := io.WriteString(w, ":keep-alive\n\n"); err != nil { + writeMu.Lock() + _, err := io.WriteString(w, ":keep-alive\n\n") + if err == nil { + flusher.Flush() + } + writeMu.Unlock() + if err != nil { return } - flusher.Flush() } } } diff --git a/pkg/bridges/subagentstream/mux.go b/pkg/bridges/subagentstream/mux.go index 33b2698..d1ba507 100644 --- a/pkg/bridges/subagentstream/mux.go +++ b/pkg/bridges/subagentstream/mux.go @@ -187,9 +187,6 @@ func AGUICustomEvent(ev a2adelegation.DelegationEvent) aguievents.Event { if ev.Error != nil { value["error"] = ev.Error } - if ev.Raw != nil { - value["raw"] = ev.Raw - } for key, val := range value { if val == "" || val == nil { delete(value, key) diff --git a/pkg/bridges/subagentstream/mux_test.go b/pkg/bridges/subagentstream/mux_test.go index 4eb1d6e..52d4697 100644 --- a/pkg/bridges/subagentstream/mux_test.go +++ b/pkg/bridges/subagentstream/mux_test.go @@ -44,6 +44,25 @@ func TestAGUICustomEventMapsDelegationFields(t *testing.T) { } } +func TestAGUICustomEventOmitsRawPayload(t *testing.T) { + t.Parallel() + ev := subagentstream.AGUICustomEvent(a2adelegation.DelegationEvent{ + RunID: "run-1", + DelegationID: "del-1", + AgentKey: "research", + Kind: a2adelegation.DelegationArtifactCreated, + Raw: map[string]any{"secret": "inline remote payload"}, + }) + custom, ok := ev.(*aguievents.CustomEvent) + if !ok { + t.Fatalf("expected CustomEvent, got %T", ev) + } + value := custom.Value.(map[string]any) + if _, ok := value["raw"]; ok { + t.Fatalf("raw payload must not be exposed to AG-UI clients: %#v", value) + } +} + func TestStreamPayloadUsesCustomPassThroughShape(t *testing.T) { t.Parallel() payload := subagentstream.StreamPayload(a2adelegation.DelegationEvent{ diff --git a/pkg/hosttools/a2adelegation/delegator.go b/pkg/hosttools/a2adelegation/delegator.go index 3d6fc81..88ca2da 100644 --- a/pkg/hosttools/a2adelegation/delegator.go +++ b/pkg/hosttools/a2adelegation/delegator.go @@ -76,6 +76,13 @@ func (d *Delegator) Delegate(ctx context.Context, req DelegationRequest) (Delega } client := d.clientFor(spec) + if client == nil { + derr := &DelegationError{Code: "configuration_error", Message: "remote agent requires AgentCardURL for default A2A execution; configure Delegator.NewClient for static AgentCard-only specs"} + d.publish(failedEvent(baseEvent, derr)) + baseResult.Status = "failed" + baseResult.Error = derr + return baseResult, derr + } card, err := client.AgentCard(ctx) if err != nil && spec.AgentCard == nil { derr := &DelegationError{Code: "agent_unavailable", Message: err.Error(), Retryable: true} @@ -355,6 +362,9 @@ func (d *Delegator) clientFor(spec RemoteAgentSpec) A2AClient { if d.NewClient != nil { return d.NewClient(spec) } + if strings.TrimSpace(spec.AgentCardURL) == "" { + return nil + } return a2aClientAdapter{Client: clienta2a.New(clienta2a.Options{ AgentCardURL: spec.AgentCardURL, Auth: spec.Auth, diff --git a/pkg/hosttools/a2adelegation/delegator_test.go b/pkg/hosttools/a2adelegation/delegator_test.go index 716af3a..d0a131e 100644 --- a/pkg/hosttools/a2adelegation/delegator_test.go +++ b/pkg/hosttools/a2adelegation/delegator_test.go @@ -349,6 +349,29 @@ func TestDelegatorRejectsInputArtifactWithoutURI(t *testing.T) { } } +func TestDelegatorStaticAgentCardRequiresCustomClient(t *testing.T) { + t.Parallel() + card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} + registry, err := NewRegistry(RemoteAgentSpec{Key: "research", AgentCard: &card}) + if err != nil { + t.Fatalf("registry: %v", err) + } + delegator := NewDelegator(registry, NewEventBus(16)) + delegator.NewID = func() string { return "del-1" } + + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this"}) + if err == nil { + t.Fatal("expected static AgentCard-only spec to require custom client") + } + var derr *DelegationError + if !errors.As(err, &derr) || derr.Code != "configuration_error" { + t.Fatalf("expected configuration_error, got result=%#v err=%T %[2]v", result, err) + } + if result.Status != "failed" { + t.Fatalf("expected failed result, got %#v", result) + } +} + func TestDelegatorMaxArtifactsLimitsFinalResultOnly(t *testing.T) { t.Parallel() card := clienta2a.AgentCard{Name: "Research", Capabilities: clienta2a.Capabilities{Streaming: false}} From 3cf440494e7c392b5d3559b70d2cd4d02fab980f Mon Sep 17 00:00:00 2001 From: buthim Date: Wed, 8 Jul 2026 17:17:55 +0800 Subject: [PATCH 6/6] Preserve delegation UI streaming without leaking runtime secrets Issue #5 requires a narrow product contract: the parent agent receives only the final MCP tool result while the host UI observes remote subagent progress on the same SSE/AG-UI stream. This hardens the run-scoped MCP sidecar, closes dangling subagent lifecycles, and adds a subprocess-only runtime secret channel so bearer tokens can be injected without entering public runtime refs or profile materialization. Constraint: agent-adaptor remains a pure SDK; delegation stays in hosttools/bridges/runtime plumbing, not concrete adapters or a second execution path Constraint: Runtime-backed MCP auth must work for local Codex and Claude parent agents without putting token values in Metadata, RuntimeServices, profile payloads, or runtime JSON Rejected: Put bearer token values in RuntimeServiceRef.Metadata | leaks through runtime events, run results, and PAPERCLIP_RUNTIME_SERVICES_JSON Rejected: Leave dangling subagent terminal handling to UI cleanup | parent terminal/cancel paths need a deterministic protocol event Confidence: high Scope-risk: moderate Directive: Keep RuntimeServiceRef.SecretEnv subprocess-only; do not copy it into public refs/reports/profile payloads or fingerprints Tested: go test -count=1 ./... Tested: go test -race -count=1 ./pkg/hosttools/a2adelegation ./pkg/bridges/subagentstream ./pkg/bridges/sse ./internal/adapterutil Tested: go vet ./... Tested: git diff --check Tested: go run ./.omx/tmp/delegation_sse_e2e --agent claude --timeout 180s Tested: go run ./.omx/tmp/delegation_sse_e2e --agent codex --timeout 180s Not-tested: SSE reconnect/replay and high-concurrency multi-delegation soak --- internal/adapterutil/helpers.go | 30 +- internal/adapterutil/helpers_test.go | 66 ++++ pkg/bridges/sse/subagent_integration_test.go | 228 ++++++++++++++ pkg/bridges/subagentstream/mux.go | 126 +++++++- pkg/bridges/subagentstream/mux_test.go | 284 +++++++++++++++++- pkg/hosttools/a2adelegation/delegator.go | 11 +- pkg/hosttools/a2adelegation/delegator_test.go | 133 +++++++- pkg/hosttools/a2adelegation/mcpserver.go | 87 +++++- pkg/hosttools/a2adelegation/types.go | 1 + runtime.go | 12 + runtime_admin_test.go | 105 ++++++- util.go | 12 + workspace_skill_types.go | 5 + 13 files changed, 1050 insertions(+), 50 deletions(-) create mode 100644 pkg/bridges/sse/subagent_integration_test.go diff --git a/internal/adapterutil/helpers.go b/internal/adapterutil/helpers.go index 45e0847..12e5df8 100644 --- a/internal/adapterutil/helpers.go +++ b/internal/adapterutil/helpers.go @@ -134,18 +134,32 @@ func ResolvedTruthyEnv(bindings []agentadaptor.EnvBinding, name string) (bool, s } func RuntimeEnvBindings(bindings []agentadaptor.EnvBinding, payload agentadaptor.RuntimePayload) ([]agentadaptor.EnvBinding, error) { - if len(payload.Ensured) == 0 { - return append([]agentadaptor.EnvBinding(nil), bindings...), nil - } - raw, err := json.Marshal(payload.Ensured) - if err != nil { - return nil, err - } out := append([]agentadaptor.EnvBinding(nil), bindings...) - out = append(out, agentadaptor.EnvBinding{Name: "PAPERCLIP_RUNTIME_SERVICES_JSON", Value: string(raw)}) + out = append(out, payload.SecretEnv...) + if len(payload.Ensured) > 0 { + raw, err := json.Marshal(runtimeServiceRefsForEnv(payload.Ensured)) + if err != nil { + return nil, err + } + out = append(out, agentadaptor.EnvBinding{Name: "PAPERCLIP_RUNTIME_SERVICES_JSON", Value: string(raw)}) + } return out, nil } +func runtimeServiceRefsForEnv(refs []agentadaptor.RuntimeServiceRef) []agentadaptor.RuntimeServiceRef { + if len(refs) == 0 { + return nil + } + out := make([]agentadaptor.RuntimeServiceRef, 0, len(refs)) + for _, ref := range refs { + clean := ref + clean.SecretEnv = nil + clean.Metadata = cloneStringMap(ref.Metadata) + out = append(out, clean) + } + return out +} + func RuntimePromptPrefix(payload agentadaptor.RuntimePayload) string { if len(payload.Ensured) == 0 { return "" diff --git a/internal/adapterutil/helpers_test.go b/internal/adapterutil/helpers_test.go index f8e787c..745c18a 100644 --- a/internal/adapterutil/helpers_test.go +++ b/internal/adapterutil/helpers_test.go @@ -1,6 +1,7 @@ package adapterutil import ( + "strings" "testing" agentadaptor "github.com/agent-dance/agent-adaptor" @@ -26,3 +27,68 @@ func TestResolvedTruthyEnvRecognizesTrueValues(t *testing.T) { t.Fatalf("unexpected truthy env result: %v %q", ok, source) } } + +func TestRuntimeEnvBindingsInjectsSecretEnvWithoutLeakingRuntimeJSON(t *testing.T) { + const secret = "sk-runtime-secret" + + env, err := RuntimeEnvBindings( + []agentadaptor.EnvBinding{{Name: "EXISTING", Value: "1"}}, + agentadaptor.RuntimePayload{ + SecretEnv: []agentadaptor.EnvBinding{{Name: "DELEGATION_TOKEN", Value: secret}}, + Ensured: []agentadaptor.RuntimeServiceRef{{ + ID: "svc-delegation", + Name: "delegation", + URL: "http://127.0.0.1:43127/mcp", + Status: agentadaptor.RuntimeServiceRunning, + Metadata: map[string]string{ + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.bearer_token_env_var": "DELEGATION_TOKEN", + }, + SecretEnv: []agentadaptor.EnvBinding{{Name: "DELEGATION_TOKEN", Value: secret}}, + }}, + }, + ) + if err != nil { + t.Fatalf("runtime env bindings: %v", err) + } + if got := lastEnvValue(env, "DELEGATION_TOKEN"); got != secret { + t.Fatalf("expected runtime secret env to reach subprocess env, got %q in %#v", got, env) + } + runtimeJSON := lastEnvValue(env, "PAPERCLIP_RUNTIME_SERVICES_JSON") + if runtimeJSON == "" { + t.Fatalf("expected runtime services JSON in env: %#v", env) + } + if strings.Contains(runtimeJSON, secret) { + t.Fatalf("runtime services JSON leaked secret: %s", runtimeJSON) + } + if !strings.Contains(runtimeJSON, "DELEGATION_TOKEN") { + t.Fatalf("runtime services JSON should retain bearer env var reference, got %s", runtimeJSON) + } +} + +func TestRuntimeEnvBindingsInjectsSecretEnvWithoutEnsuredRefs(t *testing.T) { + const secret = "sk-secret-only" + + env, err := RuntimeEnvBindings(nil, agentadaptor.RuntimePayload{ + SecretEnv: []agentadaptor.EnvBinding{{Name: "DELEGATION_TOKEN", Value: secret}}, + }) + if err != nil { + t.Fatalf("runtime env bindings: %v", err) + } + if got := lastEnvValue(env, "DELEGATION_TOKEN"); got != secret { + t.Fatalf("expected runtime secret env without ensured refs, got %q in %#v", got, env) + } + if got := lastEnvValue(env, "PAPERCLIP_RUNTIME_SERVICES_JSON"); got != "" { + t.Fatalf("did not expect runtime services JSON without ensured refs, got %q", got) + } +} + +func lastEnvValue(bindings []agentadaptor.EnvBinding, name string) string { + var out string + for _, binding := range bindings { + if binding.Name == name { + out = binding.Value + } + } + return out +} diff --git a/pkg/bridges/sse/subagent_integration_test.go b/pkg/bridges/sse/subagent_integration_test.go new file mode 100644 index 0000000..dc67aaa --- /dev/null +++ b/pkg/bridges/sse/subagent_integration_test.go @@ -0,0 +1,228 @@ +package sse_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + agentadaptor "github.com/agent-dance/agent-adaptor" + "github.com/agent-dance/agent-adaptor/memory" + "github.com/agent-dance/agent-adaptor/pkg/bridges/sse" + "github.com/agent-dance/agent-adaptor/pkg/hosttools/a2adelegation" +) + +func TestSSEHandlerAGUIRealEventBusSubagentEventsPrecedeParentFinished(t *testing.T) { + t.Parallel() + + adapter := newDelegateBoundaryAdapter() + sdk := agentadaptor.New( + agentadaptor.WithDefaultAgent(agentadaptor.Bind(adapter, nil)), + agentadaptor.WithSessionStore(memory.NewSessionStore()), + ) + bus := a2adelegation.NewEventBus(8) + srv := httptest.NewServer(sse.Handler(sdk, sse.Options{Protocol: sse.AGUI, SubagentBus: bus})) + defer srv.Close() + + body := strings.NewReader(`{ + "threadId": "t-1", + "runId": "r-1", + "messages": [{"id":"m-1","role":"user","content":"delegate this"}] + }`) + resp, err := http.Post(srv.URL, "application/json", body) + if err != nil { + t.Fatalf("post: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status: %d", resp.StatusCode) + } + + runID := adapter.waitDelegateToolActive(t) + for _, ev := range []a2adelegation.DelegationEvent{ + { + RunID: runID, + ParentToolCallID: "tool-delegate-1", + DelegationID: "del-1", + AgentKey: "research", + AgentName: "Research Agent", + Protocol: a2adelegation.ProtocolA2A, + Kind: a2adelegation.DelegationStarted, + Status: "running", + }, + { + RunID: runID, + ParentToolCallID: "tool-delegate-1", + DelegationID: "del-1", + AgentKey: "research", + Protocol: a2adelegation.ProtocolA2A, + RemoteMessageID: "msg-1", + Kind: a2adelegation.DelegationTextDelta, + Delta: "subagent text", + }, + { + RunID: runID, + ParentToolCallID: "tool-delegate-1", + DelegationID: "del-1", + AgentKey: "research", + Protocol: a2adelegation.ProtocolA2A, + RemoteArtifactID: "artifact-1", + Kind: a2adelegation.DelegationArtifactCreated, + Artifact: &a2adelegation.DelegationArtifact{ + ID: "artifact-1", + Name: "notes.md", + MediaType: "text/markdown", + }, + }, + { + RunID: runID, + ParentToolCallID: "tool-delegate-1", + DelegationID: "del-1", + AgentKey: "research", + Protocol: a2adelegation.ProtocolA2A, + Kind: a2adelegation.DelegationFinished, + Status: "completed", + Text: "done", + }, + } { + if !bus.Publish(ev) { + t.Fatalf("publish %s returned false", ev.Kind) + } + } + + adapter.finishParent() + + frames := readSSEFrames(t, resp.Body, 2*time.Second) + var types []string + customIndex := map[string]int{} + runFinishedIndex := -1 + for _, f := range frames { + if f.data == "" { + continue + } + var payload map[string]any + if err := json.Unmarshal([]byte(f.data), &payload); err != nil { + t.Fatalf("decode frame data: %v (raw=%q)", err, f.data) + } + typ, _ := payload["type"].(string) + if typ == "" { + continue + } + idx := len(types) + types = append(types, typ) + if typ == "RUN_FINISHED" { + runFinishedIndex = idx + } + if typ == "CUSTOM" { + name, _ := payload["name"].(string) + if _, exists := customIndex[name]; !exists { + customIndex[name] = idx + } + assertSubagentCustomPayload(t, payload, runID) + } + } + if runFinishedIndex < 0 { + t.Fatalf("RUN_FINISHED not observed; types=%v", types) + } + for _, name := range []string{ + string(a2adelegation.DelegationStarted), + string(a2adelegation.DelegationTextDelta), + string(a2adelegation.DelegationArtifactCreated), + string(a2adelegation.DelegationFinished), + } { + idx, ok := customIndex[name] + if !ok { + t.Fatalf("%s CUSTOM event not observed; types=%v custom=%v", name, types, customIndex) + } + if idx >= runFinishedIndex { + t.Fatalf("%s CUSTOM event index %d must precede RUN_FINISHED index %d; types=%v", name, idx, runFinishedIndex, types) + } + } +} + +func assertSubagentCustomPayload(t *testing.T, payload map[string]any, runID string) { + t.Helper() + value, _ := payload["value"].(map[string]any) + if value == nil { + t.Fatalf("CUSTOM payload missing value: %#v", payload) + } + if value["runId"] != runID || value["parentToolCallId"] != "tool-delegate-1" || + value["delegationId"] != "del-1" || value["agentKey"] != "research" { + t.Fatalf("unexpected subagent custom value: %#v", value) + } +} + +type delegateBoundaryAdapter struct { + started chan string + release chan struct{} + once sync.Once +} + +func newDelegateBoundaryAdapter() *delegateBoundaryAdapter { + return &delegateBoundaryAdapter{ + started: make(chan string, 1), + release: make(chan struct{}), + } +} + +func (a *delegateBoundaryAdapter) Descriptor() agentadaptor.DriverDescriptor { + return agentadaptor.DriverDescriptor{Type: "delegate-boundary", DisplayName: "Delegate Boundary"} +} +func (a *delegateBoundaryAdapter) ValidateConfig(any) error { return nil } +func (a *delegateBoundaryAdapter) StreamCapability() agentadaptor.StreamCapability { + return agentadaptor.StreamCapability{Native: true, TokenLevel: true, ToolCallArgs: true} +} +func (a *delegateBoundaryAdapter) Run(ctx context.Context, req agentadaptor.DriverRunRequest, sink agentadaptor.EventSink) (agentadaptor.DriverRunResult, error) { + if err := sink.EmitStream(agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunStarted, ThreadID: "t", RunID: req.RunID}); err != nil { + return agentadaptor.DriverRunResult{}, err + } + if err := sink.EmitStream(agentadaptor.StreamPayload{ + Kind: agentadaptor.StreamToolCallStart, + ToolCallID: "tool-delegate-1", + Name: a2adelegation.DelegateToolName, + }); err != nil { + return agentadaptor.DriverRunResult{}, err + } + if err := sink.EmitStream(agentadaptor.StreamPayload{ + Kind: agentadaptor.StreamToolCallArgs, + ToolCallID: "tool-delegate-1", + Delta: `{"agent":"research","objective":"collect evidence","constraints":{"stream":true}}`, + }); err != nil { + return agentadaptor.DriverRunResult{}, err + } + select { + case a.started <- req.RunID: + default: + } + + select { + case <-ctx.Done(): + return agentadaptor.DriverRunResult{}, ctx.Err() + case <-a.release: + } + if err := sink.EmitStream(agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunFinished, ThreadID: "t", RunID: req.RunID}); err != nil { + return agentadaptor.DriverRunResult{}, err + } + return agentadaptor.DriverRunResult{Output: "done", ExitCode: 0}, nil +} + +func (a *delegateBoundaryAdapter) waitDelegateToolActive(t *testing.T) string { + t.Helper() + select { + case runID := <-a.started: + return runID + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for delegate_to_agent tool call to become active") + return "" + } +} + +func (a *delegateBoundaryAdapter) finishParent() { + a.once.Do(func() { + close(a.release) + }) +} diff --git a/pkg/bridges/subagentstream/mux.go b/pkg/bridges/subagentstream/mux.go index d1ba507..f3f9b98 100644 --- a/pkg/bridges/subagentstream/mux.go +++ b/pkg/bridges/subagentstream/mux.go @@ -20,6 +20,8 @@ type MuxOptions struct { Bus EventBus } +const terminalFlushTimeout = 250 * time.Millisecond + type Event struct { ID uint64 AGUI aguievents.Event @@ -35,10 +37,8 @@ func WrapAGUI(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOption if ev.AGUI == nil { continue } - select { - case <-ctx.Done(): + if !sendAGUIWithCancelGrace(ctx, out, ev.AGUI) { return - case out <- ev.AGUI: } } }() @@ -49,19 +49,27 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < out := make(chan Event, 64) var seq atomic.Uint64 - sendAGUI := func(ev aguievents.Event) bool { - wrapped := Event{ID: seq.Add(1), AGUI: ev} + sendSubagentEvent := func(ev a2adelegation.DelegationEvent, observeContext bool) bool { + aguiEvent := AGUICustomEvent(ev) + evCopy := ev + wrapped := Event{ID: seq.Add(1), AGUI: aguiEvent, Subagent: &evCopy} + if observeContext { + select { + case <-ctx.Done(): + return false + case out <- wrapped: + return true + } + } select { - case <-ctx.Done(): - return false case out <- wrapped: return true + case <-time.After(terminalFlushTimeout): + return false } } - sendSubagent := func(ev a2adelegation.DelegationEvent) bool { - aguiEvent := AGUICustomEvent(ev) - evCopy := ev - wrapped := Event{ID: seq.Add(1), AGUI: aguiEvent, Subagent: &evCopy} + sendAGUI := func(ev aguievents.Event) bool { + wrapped := Event{ID: seq.Add(1), AGUI: ev} select { case <-ctx.Done(): return false @@ -69,6 +77,10 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < return true } } + sendSubagent := func(ev a2adelegation.DelegationEvent) bool { + return sendSubagentEvent(ev, true) + } + active := newDelegationTracker() drainSubagents := func(subagents <-chan a2adelegation.DelegationEvent) bool { for subagents != nil { select { @@ -76,6 +88,7 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < if !ok { return true } + active.Track(ev) if !sendSubagent(ev) { return false } @@ -107,6 +120,9 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < for parent != nil || subagents != nil { select { case <-ctx.Done(): + active.FlushSynthetic(a2adelegation.DelegationCancelled, "cancelled", &a2adelegation.DelegationError{Code: "parent_cancelled", Message: "parent run context cancelled"}, func(ev a2adelegation.DelegationEvent) bool { + return sendSubagentEvent(ev, false) + }) stopSubagents() cancelHandle(handle) return @@ -116,6 +132,9 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < if !drainSubagents(subagents) { return } + if !active.FlushSynthetic(a2adelegation.DelegationFailed, "failed", parentFinishedError(), sendSubagent) { + return + } stopSubagents() continue } @@ -123,6 +142,9 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < if !drainSubagents(subagents) { return } + if !active.FlushSynthetic(a2adelegation.DelegationFailed, "failed", parentFinishedError(), sendSubagent) { + return + } if !sendAGUI(ev) { return } @@ -138,6 +160,7 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < subagents = nil continue } + active.Track(ev) if !sendSubagent(ev) { return } @@ -147,6 +170,74 @@ func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) < return out } +func sendAGUIWithCancelGrace(ctx context.Context, out chan<- aguievents.Event, ev aguievents.Event) bool { + select { + case out <- ev: + return true + default: + } + if ctx.Err() != nil { + select { + case out <- ev: + return true + case <-time.After(terminalFlushTimeout): + return false + } + } + select { + case out <- ev: + return true + case <-ctx.Done(): + select { + case out <- ev: + return true + case <-time.After(terminalFlushTimeout): + return false + } + } +} + +type delegationTracker struct { + active map[string]a2adelegation.DelegationEvent + order []string +} + +func newDelegationTracker() *delegationTracker { + return &delegationTracker{active: map[string]a2adelegation.DelegationEvent{}} +} + +func (t *delegationTracker) Track(ev a2adelegation.DelegationEvent) { + if ev.DelegationID == "" { + return + } + if isTerminalDelegation(ev.Kind) { + delete(t.active, ev.DelegationID) + return + } + if _, exists := t.active[ev.DelegationID]; !exists { + t.order = append(t.order, ev.DelegationID) + } + t.active[ev.DelegationID] = ev +} + +func (t *delegationTracker) FlushSynthetic(kind a2adelegation.DelegationEventKind, status string, err *a2adelegation.DelegationError, send func(a2adelegation.DelegationEvent) bool) bool { + for _, delegationID := range t.order { + ev, ok := t.active[delegationID] + if !ok { + continue + } + ev.Kind = kind + ev.Status = status + ev.Error = err + ev.Time = time.Now() + delete(t.active, delegationID) + if !send(ev) { + return false + } + } + return true +} + func cancelHandle(handle agentadaptor.RunHandle) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() @@ -165,6 +256,19 @@ func isTerminalAGUI(ev aguievents.Event) bool { } } +func isTerminalDelegation(kind a2adelegation.DelegationEventKind) bool { + switch kind { + case a2adelegation.DelegationFinished, a2adelegation.DelegationFailed, a2adelegation.DelegationCancelled, a2adelegation.DelegationInputRequired: + return true + default: + return false + } +} + +func parentFinishedError() *a2adelegation.DelegationError { + return &a2adelegation.DelegationError{Code: "parent_finished", Message: "parent run finished before subagent terminal event"} +} + func AGUICustomEvent(ev a2adelegation.DelegationEvent) aguievents.Event { value := map[string]any{ "runId": ev.RunID, diff --git a/pkg/bridges/subagentstream/mux_test.go b/pkg/bridges/subagentstream/mux_test.go index 52d4697..df090ba 100644 --- a/pkg/bridges/subagentstream/mux_test.go +++ b/pkg/bridges/subagentstream/mux_test.go @@ -141,6 +141,124 @@ func TestWrapDrainsBufferedSubagentsBeforeParentTerminal(t *testing.T) { } } +func TestWrapSynthesizesDanglingSubagentFailureBeforeParentRunFinished(t *testing.T) { + t.Parallel() + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload, 1), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := newTrackingBus() + out := subagentstream.Wrap(context.Background(), handle, subagentstream.MuxOptions{Bus: bus}) + bus.waitSubscribed(t, "run-1") + bus.publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "active", AgentKey: "research", Kind: a2adelegation.DelegationStarted}) + bus.publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "done", AgentKey: "review", Kind: a2adelegation.DelegationStarted}) + bus.publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "done", AgentKey: "review", Kind: a2adelegation.DelegationFinished, Status: "completed"}) + + handle.stream <- agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunFinished, ThreadID: "thread-1", RunID: "run-1"} + close(handle.stream) + close(handle.done) + events := collectAllMuxEvents(t, out) + + assertLastAGUIType(t, events, aguievents.EventTypeRunFinished) + assertSyntheticSubagentTerminalBeforeParent(t, events, "active", a2adelegation.DelegationFailed, "parent_finished") + assertNoExtraTerminal(t, events, "done", a2adelegation.DelegationFailed) +} + +func TestWrapSynthesizesDanglingSubagentFailureBeforeSynthesizedParentTerminalOnStreamClose(t *testing.T) { + t.Parallel() + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := newTrackingBus() + out := subagentstream.Wrap(context.Background(), handle, subagentstream.MuxOptions{Bus: bus}) + bus.waitSubscribed(t, "run-1") + bus.publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "active", AgentKey: "research", Kind: a2adelegation.DelegationStatus, Status: "working"}) + + close(handle.stream) + close(handle.done) + events := collectAllMuxEvents(t, out) + + assertLastAGUIType(t, events, aguievents.EventTypeRunFinished) + assertSyntheticSubagentTerminalBeforeParent(t, events, "active", a2adelegation.DelegationFailed, "parent_finished") +} + +func TestWrapSynthesizesDanglingSubagentFailureBeforeParentRunError(t *testing.T) { + t.Parallel() + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload, 1), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := newTrackingBus() + out := subagentstream.Wrap(context.Background(), handle, subagentstream.MuxOptions{Bus: bus}) + bus.waitSubscribed(t, "run-1") + bus.publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "active", AgentKey: "research", Kind: a2adelegation.DelegationStarted}) + + handle.stream <- agentadaptor.StreamPayload{Kind: agentadaptor.StreamRunError, ThreadID: "thread-1", RunID: "run-1", Error: &agentadaptor.RunFailure{Message: "boom", Code: "parent.error"}} + close(handle.stream) + close(handle.done) + events := collectAllMuxEvents(t, out) + + assertLastAGUIType(t, events, aguievents.EventTypeRunError) + assertSyntheticSubagentTerminalBeforeParent(t, events, "active", a2adelegation.DelegationFailed, "parent_finished") +} + +func TestWrapSynthesizesDanglingSubagentCancelledOnContextCancel(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := newTrackingBus() + out := subagentstream.Wrap(ctx, handle, subagentstream.MuxOptions{Bus: bus}) + bus.waitSubscribed(t, "run-1") + bus.publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "active", AgentKey: "research", Kind: a2adelegation.DelegationStarted}) + + events := collectMuxEvents(t, out, 1) + cancel() + events = append(events, collectAllMuxEvents(t, out)...) + + assertSubagentTerminal(t, events, "active", a2adelegation.DelegationCancelled, "parent_cancelled") + if !handle.sawLiveCancelContext() { + t.Fatalf("expected at least one Cancel call with a live context, got errs=%v", handle.cancelErrs()) + } +} + +func TestWrapAGUIForwardsDanglingSubagentCancelledOnContextCancel(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + handle := &fakeHandle{ + stream: make(chan agentadaptor.StreamPayload), + events: make(chan agentadaptor.RunEvent), + done: make(chan struct{}), + runID: "run-1", + runResult: agentadaptor.RunResult{RunID: "run-1"}, + } + bus := newTrackingBus() + out := subagentstream.WrapAGUI(ctx, handle, subagentstream.MuxOptions{Bus: bus}) + bus.waitSubscribed(t, "run-1") + bus.publish(a2adelegation.DelegationEvent{RunID: "run-1", DelegationID: "active", AgentKey: "research", Kind: a2adelegation.DelegationStarted}) + + events := collectAGUIEvents(t, out, 1) + cancel() + events = append(events, collectAllAGUIEvents(t, out)...) + + assertAGUICustomTerminal(t, events, "active", a2adelegation.DelegationCancelled, "parent_cancelled") +} + func TestWrapAGUIExitsWhenContextCanceledAndConsumerStops(t *testing.T) { t.Parallel() ctx, cancel := context.WithCancel(context.Background()) @@ -236,22 +354,156 @@ func collectMuxEvents(t *testing.T, ch <-chan subagentstream.Event, want int) [] t.Fatalf("mux closed after %d events, wanted %d", len(out), want) } out = append(out, ev) - case <-time.After(time.Second): + case <-time.After(3 * time.Second): t.Fatalf("timeout waiting for mux event %d/%d", len(out), want) } } return out } +func collectAllMuxEvents(t *testing.T, ch <-chan subagentstream.Event) []subagentstream.Event { + t.Helper() + var out []subagentstream.Event + for { + select { + case ev, ok := <-ch: + if !ok { + return out + } + out = append(out, ev) + case <-time.After(3 * time.Second): + t.Fatalf("timeout waiting for mux to close after %d events", len(out)) + } + } +} + +func assertLastAGUIType(t *testing.T, events []subagentstream.Event, want aguievents.EventType) { + t.Helper() + if len(events) == 0 { + t.Fatal("expected mux events") + } + if events[len(events)-1].AGUI.Type() != want { + t.Fatalf("last AG-UI event: got %s want %s; events=%#v", events[len(events)-1].AGUI.Type(), want, events) + } +} + +func assertSyntheticSubagentTerminalBeforeParent(t *testing.T, events []subagentstream.Event, delegationID string, kind a2adelegation.DelegationEventKind, code string) { + t.Helper() + parentTerminal := len(events) - 1 + seen := 0 + for i, ev := range events { + if ev.Subagent == nil || ev.Subagent.DelegationID != delegationID || ev.Subagent.Kind != kind { + continue + } + seen++ + if i >= parentTerminal { + t.Fatalf("synthetic terminal for %q was not before parent terminal: index=%d parent=%d", delegationID, i, parentTerminal) + } + if ev.Subagent.Error == nil || ev.Subagent.Error.Code != code { + t.Fatalf("synthetic terminal error: got %#v want code %q", ev.Subagent.Error, code) + } + } + if seen != 1 { + t.Fatalf("synthetic terminal count for %q: got %d want 1; events=%#v", delegationID, seen, events) + } +} + +func assertSubagentTerminal(t *testing.T, events []subagentstream.Event, delegationID string, kind a2adelegation.DelegationEventKind, code string) { + t.Helper() + seen := 0 + for _, ev := range events { + if ev.Subagent == nil || ev.Subagent.DelegationID != delegationID || ev.Subagent.Kind != kind { + continue + } + seen++ + if ev.Subagent.Error == nil || ev.Subagent.Error.Code != code { + t.Fatalf("subagent terminal error: got %#v want code %q", ev.Subagent.Error, code) + } + } + if seen != 1 { + t.Fatalf("subagent terminal count for %q: got %d want 1; events=%#v", delegationID, seen, events) + } +} + +func collectAGUIEvents(t *testing.T, ch <-chan aguievents.Event, want int) []aguievents.Event { + t.Helper() + out := make([]aguievents.Event, 0, want) + for len(out) < want { + select { + case ev, ok := <-ch: + if !ok { + t.Fatalf("AG-UI mux closed after %d events, wanted %d", len(out), want) + } + out = append(out, ev) + case <-time.After(3 * time.Second): + t.Fatalf("timeout waiting for AG-UI event %d/%d", len(out), want) + } + } + return out +} + +func collectAllAGUIEvents(t *testing.T, ch <-chan aguievents.Event) []aguievents.Event { + t.Helper() + var out []aguievents.Event + for { + select { + case ev, ok := <-ch: + if !ok { + return out + } + out = append(out, ev) + case <-time.After(3 * time.Second): + t.Fatalf("timeout waiting for AG-UI mux to close after %d events", len(out)) + } + } +} + +func assertAGUICustomTerminal(t *testing.T, events []aguievents.Event, delegationID string, kind a2adelegation.DelegationEventKind, code string) { + t.Helper() + seen := 0 + for _, ev := range events { + custom, ok := ev.(*aguievents.CustomEvent) + if !ok || custom.Name != string(kind) { + continue + } + value, ok := custom.Value.(map[string]any) + if !ok { + t.Fatalf("custom event value: got %T", custom.Value) + } + if value["delegationId"] != delegationID { + continue + } + seen++ + derr, ok := value["error"].(*a2adelegation.DelegationError) + if !ok || derr.Code != code { + t.Fatalf("custom terminal error: got %#v want code %q", value["error"], code) + } + } + if seen != 1 { + t.Fatalf("custom terminal count for %q: got %d want 1; events=%#v", delegationID, seen, events) + } +} + +func assertNoExtraTerminal(t *testing.T, events []subagentstream.Event, delegationID string, kind a2adelegation.DelegationEventKind) { + t.Helper() + for _, ev := range events { + if ev.Subagent != nil && ev.Subagent.DelegationID == delegationID && ev.Subagent.Kind == kind { + t.Fatalf("delegation %q got duplicate terminal %s in events=%#v", delegationID, kind, events) + } + } +} + type trackingBus struct { - mu sync.Mutex - runID string - ctxDone <-chan struct{} - events chan a2adelegation.DelegationEvent + mu sync.Mutex + readyOnce sync.Once + runID string + ctxDone <-chan struct{} + ready chan struct{} + events chan a2adelegation.DelegationEvent } func newTrackingBus() *trackingBus { - return &trackingBus{events: make(chan a2adelegation.DelegationEvent)} + return &trackingBus{ready: make(chan struct{}), events: make(chan a2adelegation.DelegationEvent, 16)} } func (b *trackingBus) SubscribeRun(ctx context.Context, runID string) <-chan a2adelegation.DelegationEvent { @@ -259,9 +511,29 @@ func (b *trackingBus) SubscribeRun(ctx context.Context, runID string) <-chan a2a b.runID = runID b.ctxDone = ctx.Done() b.mu.Unlock() + b.readyOnce.Do(func() { close(b.ready) }) return b.events } +func (b *trackingBus) waitSubscribed(t *testing.T, wantRunID string) { + t.Helper() + select { + case <-b.ready: + case <-time.After(3 * time.Second): + t.Fatal("timeout waiting for subagent subscription") + } + b.mu.Lock() + runID := b.runID + b.mu.Unlock() + if runID != wantRunID { + t.Fatalf("subscribed runID: got %q want %q", runID, wantRunID) + } +} + +func (b *trackingBus) publish(ev a2adelegation.DelegationEvent) { + b.events <- ev +} + func (b *trackingBus) assertCanceled(t *testing.T, wantRunID string) { t.Helper() b.mu.Lock() diff --git a/pkg/hosttools/a2adelegation/delegator.go b/pkg/hosttools/a2adelegation/delegator.go index 88ca2da..1ed2049 100644 --- a/pkg/hosttools/a2adelegation/delegator.go +++ b/pkg/hosttools/a2adelegation/delegator.go @@ -114,6 +114,7 @@ func (d *Delegator) Delegate(ctx context.Context, req DelegationRequest) (Delega Message: message, Tenant: spec.Tenant, AcceptedOutputModes: spec.AcceptedOutputModes, + HistoryLength: req.HistoryLength, Metadata: cloneAnyMap(req.Metadata), } if req.Tenant != "" { @@ -185,7 +186,7 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe return cancelResult() } if lastTaskID != "" { - if recovered, ok := d.recoverTask(ctx, client, lastTaskID, send.Tenant); ok { + if recovered, ok := d.recoverTask(ctx, client, lastTaskID, send.Tenant, send.HistoryLength); ok { lastTask = recovered for _, ev := range mapper.taskEvents(recovered) { d.publish(ev) @@ -229,7 +230,7 @@ func (d *Delegator) delegateStreaming(ctx context.Context, client A2AClient, spe return d.finishTask(baseEvent, baseResult, *event.Task, spec.Policy, maxArtifacts) } if event.Status != nil && executionFinalState(event.Status.State) { - task, ok := d.recoverTask(ctx, client, event.TaskID, send.Tenant) + task, ok := d.recoverTask(ctx, client, event.TaskID, send.Tenant, send.HistoryLength) if ok { return d.finishTask(baseEvent, baseResult, task, spec.Policy, maxArtifacts) } @@ -304,7 +305,7 @@ func (d *Delegator) delegatePolling(ctx context.Context, client A2AClient, spec return baseResult, derr case <-ticker.C: } - task, err = client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: task.ID, Tenant: send.Tenant}) + task, err = client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: task.ID, Tenant: send.Tenant, HistoryLength: send.HistoryLength}) if err != nil { continue } @@ -322,11 +323,11 @@ func (d *Delegator) delegatePolling(ctx context.Context, client A2AClient, spec return baseResult, derr } -func (d *Delegator) recoverTask(ctx context.Context, client A2AClient, taskID, tenant string) (clienta2a.Task, bool) { +func (d *Delegator) recoverTask(ctx context.Context, client A2AClient, taskID, tenant string, historyLength *int) (clienta2a.Task, bool) { if taskID == "" { return clienta2a.Task{}, false } - task, err := client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: taskID, Tenant: tenant}) + task, err := client.GetTask(ctx, clienta2a.GetTaskRequest{TaskID: taskID, Tenant: tenant, HistoryLength: historyLength}) return task, err == nil && executionFinalState(task.Status.State) } diff --git a/pkg/hosttools/a2adelegation/delegator_test.go b/pkg/hosttools/a2adelegation/delegator_test.go index d0a131e..e6ca198 100644 --- a/pkg/hosttools/a2adelegation/delegator_test.go +++ b/pkg/hosttools/a2adelegation/delegator_test.go @@ -71,6 +71,34 @@ func TestParseToolInputTracksExplicitMaxArtifacts(t *testing.T) { } } +func TestParseToolInputRejectsInvalidTimeouts(t *testing.T) { + t.Parallel() + cases := []string{ + `{"agent":"research","objective":"do work","constraints":{"timeout_seconds":0}}`, + `{"agent":"research","objective":"do work","constraints":{"timeout_seconds":-1}}`, + `{"agent":"research","objective":"do work","constraints":{"timeout_seconds":9223372037}}`, + } + for _, raw := range cases { + if _, err := ParseToolInput([]byte(raw)); err == nil { + t.Fatalf("expected invalid timeout to fail for %s", raw) + } + } +} + +func TestParseToolInputTracksHistoryLength(t *testing.T) { + t.Parallel() + input, err := ParseToolInput([]byte(`{"agent":"research","objective":"do work","constraints":{"history_length":0}}`)) + if err != nil { + t.Fatalf("parse explicit history_length: %v", err) + } + if !input.Constraints.HistoryLengthSet || input.Constraints.HistoryLength != 0 { + t.Fatalf("expected explicit history_length=0, got %#v", input.Constraints) + } + if _, err := ParseToolInput([]byte(`{"agent":"research","objective":"do work","constraints":{"history_length":-1}}`)); err == nil { + t.Fatal("expected negative history_length to fail") + } +} + func TestToolSchemaAllowsOnlyRegistryKeyObjectiveInputAndConstraints(t *testing.T) { t.Parallel() schema := ToolSchema() @@ -188,8 +216,9 @@ func TestDelegatorPollingHappyPathEmitsEventsAndStructuredResult(t *testing.T) { delegator := NewDelegator(registry, bus) delegator.NewClient = func(RemoteAgentSpec) A2AClient { return client } delegator.NewID = func() string { return "del-1" } + historyLength := 7 - result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this"}) + result, err := delegator.Delegate(context.Background(), DelegationRequest{RunID: "run-1", Agent: "research", Objective: "research this", HistoryLength: &historyLength}) if err != nil { t.Fatalf("delegate: %v", err) } @@ -199,6 +228,12 @@ func TestDelegatorPollingHappyPathEmitsEventsAndStructuredResult(t *testing.T) { if client.lastSend.Tenant != "tenant-a" || client.lastSend.Message.Parts[0].Text != "research this" { t.Fatalf("unexpected send request: %#v", client.lastSend) } + if client.lastSend.HistoryLength == nil || *client.lastSend.HistoryLength != historyLength { + t.Fatalf("expected send history length %d, got %#v", historyLength, client.lastSend.HistoryLength) + } + if client.lastGet.HistoryLength == nil || *client.lastGet.HistoryLength != historyLength { + t.Fatalf("expected get history length %d, got %#v", historyLength, client.lastGet.HistoryLength) + } replayed := drainBus(t, bus, "run-1", 5) if replayed[0].Kind != DelegationStarted || replayed[len(replayed)-1].Kind != DelegationFinished { t.Fatalf("unexpected bus replay: %#v", replayed) @@ -818,7 +853,7 @@ func TestMCPServerDelegatesToolCallAndReturnsStructuredResult(t *testing.T) { server := httptest.NewServer(NewMCPServer(delegator, MCPServerOptions{RunID: "run-1", BearerToken: "token"}).Handler()) defer server.Close() - rpcBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"delegate_to_agent","arguments":{"agent":"research","objective":"research this","constraints":{"stream":false}}}}` + rpcBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"delegate_to_agent","arguments":{"agent":"research","objective":"research this","constraints":{"stream":false,"history_length":3}}}}` req, err := http.NewRequest(http.MethodPost, server.URL, strings.NewReader(rpcBody)) if err != nil { t.Fatalf("request: %v", err) @@ -853,6 +888,9 @@ func TestMCPServerDelegatesToolCallAndReturnsStructuredResult(t *testing.T) { if result.DelegationID != "del-1" || result.RemoteTaskID != "task-1" || result.Status != "completed" || result.Summary != "done" { t.Fatalf("unexpected delegation result: %#v", result) } + if client.lastSend.HistoryLength == nil || *client.lastSend.HistoryLength != 3 { + t.Fatalf("expected history_length to flow into send request, got %#v", client.lastSend.HistoryLength) + } } func TestMCPServerRejectsMissingBearerToken(t *testing.T) { @@ -869,6 +907,97 @@ func TestMCPServerRejectsMissingBearerToken(t *testing.T) { } } +func TestMCPServerRejectsEmptyBearerTokenByDefault(t *testing.T) { + t.Parallel() + server := httptest.NewServer(NewMCPServer(NewDelegator(nil, nil), MCPServerOptions{RunID: "run-1"}).Handler()) + defer server.Close() + resp, err := http.Post(server.URL, "application/json", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + if err != nil { + t.Fatalf("post rpc: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("expected missing server bearer token to fail closed, got %d", resp.StatusCode) + } +} + +func TestMCPServerAllowsUnauthenticatedLoopbackWhenExplicit(t *testing.T) { + t.Parallel() + server := httptest.NewServer(NewMCPServer(NewDelegator(nil, nil), MCPServerOptions{RunID: "run-1", AllowUnauthenticatedLoopbackForTest: true}).Handler()) + defer server.Close() + resp, err := http.Post(server.URL, "application/json", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + if err != nil { + t.Fatalf("post rpc: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("expected explicit loopback test mode to allow request, got %d", resp.StatusCode) + } +} + +func TestMCPServerToolCallRequiresRunContext(t *testing.T) { + t.Parallel() + server := httptest.NewServer(NewMCPServer(NewDelegator(nil, nil), MCPServerOptions{BearerToken: "token"}).Handler()) + defer server.Close() + result, isError := callDelegateTool(t, server.URL, "token", `{"agent":"research","objective":"research this"}`) + if !isError { + t.Fatalf("expected missing run context to be an MCP tool error: %#v", result) + } + if result.Error == nil || result.Error.Code != "configuration_error" || !strings.Contains(result.Error.Message, "run context") { + t.Fatalf("expected run-context configuration_error, got %#v", result) + } +} + +func TestMCPServerPreservesDelegatorErrorDetails(t *testing.T) { + t.Parallel() + server := httptest.NewServer(NewMCPServer(NewDelegator(nil, nil), MCPServerOptions{RunID: "run-1", BearerToken: "token"}).Handler()) + defer server.Close() + result, isError := callDelegateTool(t, server.URL, "token", `{"agent":"research","objective":"research this"}`) + if !isError { + t.Fatalf("expected delegator failure to be an MCP tool error: %#v", result) + } + if result.Status != "failed" || result.Error == nil || result.Error.Code != "configuration_error" { + t.Fatalf("expected structured configuration_error to be preserved, got %#v", result) + } +} + +func callDelegateTool(t *testing.T, url, token, arguments string) (DelegationResult, bool) { + t.Helper() + rpcBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"delegate_to_agent","arguments":` + arguments + `}}` + req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(rpcBody)) + if err != nil { + t.Fatalf("request: %v", err) + } + req.Header.Set("Authorization", "Bearer "+token) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("post rpc: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status: %d", resp.StatusCode) + } + var envelope struct { + Result struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + IsError bool `json:"isError"` + } `json:"result"` + } + if err := json.NewDecoder(resp.Body).Decode(&envelope); err != nil { + t.Fatalf("decode response: %v", err) + } + if len(envelope.Result.Content) != 1 { + t.Fatalf("expected one content item, got %#v", envelope.Result) + } + var result DelegationResult + if err := json.Unmarshal([]byte(envelope.Result.Content[0].Text), &result); err != nil { + t.Fatalf("decode delegation result: %v", err) + } + return result, envelope.Result.IsError +} + type fakeA2AClient struct { card clienta2a.AgentCard cardErr error diff --git a/pkg/hosttools/a2adelegation/mcpserver.go b/pkg/hosttools/a2adelegation/mcpserver.go index 953b4f5..a1fac99 100644 --- a/pkg/hosttools/a2adelegation/mcpserver.go +++ b/pkg/hosttools/a2adelegation/mcpserver.go @@ -2,14 +2,18 @@ package a2adelegation import ( "bytes" + "crypto/subtle" "encoding/json" "errors" "fmt" + "net" "net/http" "strings" "time" ) +const maxToolTimeoutSeconds = int64(1<<63-1) / int64(time.Second) + type ToolInput struct { Agent string `json:"agent"` Objective string `json:"objective"` @@ -24,10 +28,13 @@ type ToolInputBody struct { } type ToolConstraints struct { - TimeoutSeconds int `json:"timeout_seconds,omitempty"` - Stream bool `json:"stream,omitempty"` - MaxArtifacts int `json:"max_artifacts,omitempty"` - MaxArtifactsSet bool `json:"-"` + TimeoutSeconds int `json:"timeout_seconds,omitempty"` + TimeoutSecondsSet bool `json:"-"` + Stream bool `json:"stream,omitempty"` + MaxArtifacts int `json:"max_artifacts,omitempty"` + MaxArtifactsSet bool `json:"-"` + HistoryLength int `json:"history_length,omitempty"` + HistoryLengthSet bool `json:"-"` } func ToolSchema() map[string]any { @@ -62,6 +69,7 @@ func ToolSchema() map[string]any { "timeout_seconds": map[string]any{"type": "integer", "minimum": 1}, "stream": map[string]any{"type": "boolean"}, "max_artifacts": map[string]any{"type": "integer", "minimum": 0}, + "history_length": map[string]any{"type": "integer", "minimum": 0}, }, "additionalProperties": false, }, @@ -100,14 +108,31 @@ func ParseToolInput(raw []byte) (ToolInput, error) { if rawConstraints, ok := envelope["constraints"]; ok { var constraints map[string]json.RawMessage if err := json.Unmarshal(rawConstraints, &constraints); err == nil { + if _, ok := constraints["timeout_seconds"]; ok { + input.Constraints.TimeoutSecondsSet = true + } if _, ok := constraints["max_artifacts"]; ok { input.Constraints.MaxArtifactsSet = true } + if _, ok := constraints["history_length"]; ok { + input.Constraints.HistoryLengthSet = true + } + } + } + if input.Constraints.TimeoutSecondsSet { + if input.Constraints.TimeoutSeconds <= 0 { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "timeout_seconds must be positive"} + } + if int64(input.Constraints.TimeoutSeconds) > maxToolTimeoutSeconds { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "timeout_seconds exceeds maximum duration"} } } if input.Constraints.MaxArtifactsSet && input.Constraints.MaxArtifacts < 0 { return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "max_artifacts must be non-negative"} } + if input.Constraints.HistoryLengthSet && input.Constraints.HistoryLength < 0 { + return ToolInput{}, &DelegationError{Code: "invalid_tool_input", Message: "history_length must be non-negative"} + } return input, nil } @@ -116,6 +141,11 @@ type MCPServerOptions struct { ParentToolCallID string Tenant string BearerToken string + + // AllowUnauthenticatedLoopbackForTest permits an otherwise unprotected + // loopback-only server for tests and local probes. Production HTTP sidecars + // should always use a per-run bearer token. + AllowUnauthenticatedLoopbackForTest bool } type MCPServer struct { @@ -154,17 +184,32 @@ func (s *MCPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { } func (s *MCPServer) authorize(r *http.Request) error { - if s == nil || s.Options.BearerToken == "" { - return nil + if s == nil { + return errors.New("delegation MCP server is not configured") + } + if s.Options.BearerToken == "" { + if s.Options.AllowUnauthenticatedLoopbackForTest && isLoopbackRequest(r) { + return nil + } + return errors.New("delegation MCP bearer token is required") } got := strings.TrimSpace(r.Header.Get("Authorization")) want := "Bearer " + s.Options.BearerToken - if got != want { + if subtle.ConstantTimeCompare([]byte(got), []byte(want)) != 1 { return errors.New("unauthorized") } return nil } +func isLoopbackRequest(r *http.Request) bool { + host, _, err := net.SplitHostPort(r.RemoteAddr) + if err != nil { + host = r.RemoteAddr + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} + func (s *MCPServer) handle(r *http.Request, req rpcRequest) map[string]any { switch req.Method { case "initialize": @@ -206,8 +251,13 @@ func (s *MCPServer) handleToolCall(r *http.Request, req rpcRequest) map[string]a if err != nil { return toolResult(req.ID, DelegationResult{Status: "failed", Error: delegationErr(err)}, true) } + runID := strings.TrimSpace(s.Options.RunID) + if runID == "" { + derr := &DelegationError{Code: "configuration_error", Message: "run context is required for delegation MCP tool"} + return toolResult(req.ID, DelegationResult{Agent: input.Agent, RemoteProtocol: ProtocolA2A, Status: "failed", Error: derr}, true) + } request := DelegationRequest{ - RunID: s.Options.RunID, + RunID: runID, ParentToolCallID: s.Options.ParentToolCallID, Agent: input.Agent, Objective: input.Objective, @@ -220,13 +270,32 @@ func (s *MCPServer) handleToolCall(r *http.Request, req rpcRequest) map[string]a if input.Constraints.MaxArtifactsSet { request.MaxArtifacts = &input.Constraints.MaxArtifacts } - if input.Constraints.TimeoutSeconds > 0 { + if input.Constraints.HistoryLengthSet { + request.HistoryLength = &input.Constraints.HistoryLength + } + if input.Constraints.TimeoutSecondsSet { request.Timeout = time.Duration(input.Constraints.TimeoutSeconds) * time.Second } out, err := s.Delegator.Delegate(r.Context(), request) + if err != nil { + out = ensureDelegationError(out, err) + } return toolResult(req.ID, out, err != nil) } +func ensureDelegationError(out DelegationResult, err error) DelegationResult { + if out.Status == "" { + out.Status = "failed" + } + if out.RemoteProtocol == "" { + out.RemoteProtocol = ProtocolA2A + } + if out.Error == nil { + out.Error = delegationErr(err) + } + return out +} + func delegationErr(err error) *DelegationError { var derr *DelegationError if errors.As(err, &derr) { diff --git a/pkg/hosttools/a2adelegation/types.go b/pkg/hosttools/a2adelegation/types.go index 061552b..bb907cf 100644 --- a/pkg/hosttools/a2adelegation/types.go +++ b/pkg/hosttools/a2adelegation/types.go @@ -125,6 +125,7 @@ type DelegationRequest struct { Context string Artifacts []InputArtifact MaxArtifacts *int + HistoryLength *int Timeout time.Duration Stream bool Tenant string diff --git a/runtime.go b/runtime.go index 353d282..d3d778d 100644 --- a/runtime.go +++ b/runtime.go @@ -56,10 +56,22 @@ func (s *sdkImpl) prepareRuntime( if err != nil { return RuntimePayload{}, err } + payload.SecretEnv = collectRuntimeSecretEnv(ensured) payload.Ensured = normalizeRuntimeServiceRefs(payload.Requested, ensured, identity) return payload, nil } +func collectRuntimeSecretEnv(refs []RuntimeServiceRef) []EnvBinding { + if len(refs) == 0 { + return nil + } + var out []EnvBinding + for _, ref := range refs { + out = append(out, cloneEnvBindings(ref.SecretEnv)...) + } + return out +} + func runtimeReportsFromRefs(refs []RuntimeServiceRef, owner AgentIdentity) []RuntimeServiceReport { if len(refs) == 0 { return nil diff --git a/runtime_admin_test.go b/runtime_admin_test.go index 78a8e48..aaa399d 100644 --- a/runtime_admin_test.go +++ b/runtime_admin_test.go @@ -2,6 +2,8 @@ package agentadaptor_test import ( "context" + "encoding/json" + "strings" "testing" agentadaptor "github.com/agent-dance/agent-adaptor" @@ -192,6 +194,8 @@ func TestRuntimeServicesFlowThroughSingleExecutionPath(t *testing.T) { } func TestRuntimeServiceMetadataInjectsMCPIntoDriverAndProfile(t *testing.T) { + const secret = "sk-runtime-secret" + driver := &runtimeAdminDriver{ mcpCapability: agentadaptor.MCPCapability{Supported: true, HTTP: true}, } @@ -201,15 +205,16 @@ func TestRuntimeServiceMetadataInjectsMCPIntoDriverAndProfile(t *testing.T) { Name: "a2a-delegation", URL: "http://127.0.0.1:43127/mcp", Metadata: map[string]string{ - "agentadaptor.mcp.enabled": "true", - "agentadaptor.mcp.key": "delegate-a2a", - "agentadaptor.mcp.transport": "http", - "agentadaptor.mcp.headers_json": `{"X-Run-Token":"env:DELEGATION_TOKEN"}`, - "agentadaptor.mcp.bearer_token_env_var": "DELEGATION_TOKEN", - "agentadaptor.mcp.required": "true", - "agentadaptor.mcp.required_reason": "visual A2A subagent delegation", - "delegation.secret_should_not_be_promoted": "sk-test-secret", + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.key": "delegate-a2a", + "agentadaptor.mcp.transport": "http", + "agentadaptor.mcp.headers_json": `{"X-Run-Token":"env:DELEGATION_TOKEN"}`, + "agentadaptor.mcp.bearer_token_env_var": "DELEGATION_TOKEN", + "agentadaptor.mcp.required": "true", + "agentadaptor.mcp.required_reason": "visual A2A subagent delegation", + "delegation.visibility": "public-metadata", }, + SecretEnv: []agentadaptor.EnvBinding{{Name: "DELEGATION_TOKEN", Value: secret}}, }, } sdk := agentadaptor.New( @@ -220,7 +225,8 @@ func TestRuntimeServiceMetadataInjectsMCPIntoDriverAndProfile(t *testing.T) { agentadaptor.WithRuntimeServiceManager(runtimeManager), ) - if _, err := sdk.Run(context.Background(), "use delegation"); err != nil { + result, err := sdk.Run(context.Background(), "use delegation") + if err != nil { t.Fatalf("run: %v", err) } if len(driver.lastReq.MCP.Servers) != 2 { @@ -241,6 +247,16 @@ func TestRuntimeServiceMetadataInjectsMCPIntoDriverAndProfile(t *testing.T) { if injected.Headers["X-Run-Token"] != "env:DELEGATION_TOKEN" || injected.BearerTokenEnvVar != "DELEGATION_TOKEN" { t.Fatalf("expected run-scoped auth references, got %#v", injected) } + if !hasEnvBinding(driver.lastReq.Runtime.SecretEnv, "DELEGATION_TOKEN", secret) { + t.Fatalf("expected secret env in adapter runtime payload, got %#v", driver.lastReq.Runtime.SecretEnv) + } + if len(driver.lastReq.Runtime.Ensured) != 1 || len(driver.lastReq.Runtime.Ensured[0].SecretEnv) != 0 { + t.Fatalf("expected sanitized ensured runtime refs, got %#v", driver.lastReq.Runtime.Ensured) + } + assertJSONDoesNotContain(t, "driver runtime ensured refs", driver.lastReq.Runtime.Ensured, secret) + assertJSONDoesNotContain(t, "run result runtime services", result.RuntimeServices, secret) + assertJSONDoesNotContain(t, "profile payload", driver.lastReq.ProfilePayload, secret) + assertJSONDoesNotContain(t, "mcp payload", driver.lastReq.MCP, secret) if !injected.Required || injected.RequiredReason != "visual A2A subagent delegation" { t.Fatalf("expected required runtime MCP server, got %#v", injected) } @@ -252,6 +268,57 @@ func TestRuntimeServiceMetadataInjectsMCPIntoDriverAndProfile(t *testing.T) { } } +func TestRuntimeServiceSecretEnvDoesNotChangeProfileFingerprint(t *testing.T) { + driver := &runtimeAdminDriver{ + mcpCapability: agentadaptor.MCPCapability{Supported: true, HTTP: true}, + } + runtimeManager := &runtimeMCPManager{ + ref: runtimeMCPRefWithSecret("sk-first-secret"), + } + sdk := agentadaptor.New( + agentadaptor.WithDefaultAgent(agentadaptor.Bind(driver, fakeConfig{Label: "runtime"}, + agentadaptor.WithDefaultRuntimeServices(agentadaptor.RuntimeServiceSpec{Name: "a2a-delegation"}), + )), + agentadaptor.WithRuntimeServiceManager(runtimeManager), + ) + + if _, err := sdk.Run(context.Background(), "first"); err != nil { + t.Fatalf("first run: %v", err) + } + firstProfile := driver.lastReq.ProfilePayload.Fingerprint + firstMCP := driver.lastReq.MCP.Fingerprint + + runtimeManager.ref = runtimeMCPRefWithSecret("sk-second-secret") + if _, err := sdk.Run(context.Background(), "second"); err != nil { + t.Fatalf("second run: %v", err) + } + if driver.lastReq.ProfilePayload.Fingerprint != firstProfile { + t.Fatalf("secret rotation changed profile fingerprint: %q -> %q", firstProfile, driver.lastReq.ProfilePayload.Fingerprint) + } + if driver.lastReq.MCP.Fingerprint != firstMCP { + t.Fatalf("secret rotation changed MCP fingerprint: %q -> %q", firstMCP, driver.lastReq.MCP.Fingerprint) + } + if !hasEnvBinding(driver.lastReq.Runtime.SecretEnv, "DELEGATION_TOKEN", "sk-second-secret") { + t.Fatalf("expected rotated secret env in runtime payload, got %#v", driver.lastReq.Runtime.SecretEnv) + } + assertJSONDoesNotContain(t, "profile payload", driver.lastReq.ProfilePayload, "sk-second-secret") +} + +func runtimeMCPRefWithSecret(secret string) agentadaptor.RuntimeServiceRef { + return agentadaptor.RuntimeServiceRef{ + ID: "svc-delegation", + Name: "a2a-delegation", + URL: "http://127.0.0.1:43127/mcp", + Metadata: map[string]string{ + "agentadaptor.mcp.enabled": "true", + "agentadaptor.mcp.key": "delegate-a2a", + "agentadaptor.mcp.transport": "http", + "agentadaptor.mcp.bearer_token_env_var": "DELEGATION_TOKEN", + }, + SecretEnv: []agentadaptor.EnvBinding{{Name: "DELEGATION_TOKEN", Value: secret}}, + } +} + func TestRuntimeServiceMCPMetadataRejectsMalformedJSON(t *testing.T) { driver := &runtimeAdminDriver{ mcpCapability: agentadaptor.MCPCapability{Supported: true, HTTP: true}, @@ -352,6 +419,26 @@ func (m *runtimeMCPManager) ReleaseByRun(_ context.Context, _ string) error { re func (m *runtimeMCPManager) ReleaseByLabels(_ context.Context, _ map[string]string) error { return nil } +func hasEnvBinding(bindings []agentadaptor.EnvBinding, name, value string) bool { + for _, binding := range bindings { + if binding.Name == name && binding.Value == value { + return true + } + } + return false +} + +func assertJSONDoesNotContain(t *testing.T, label string, value any, forbidden string) { + t.Helper() + raw, err := json.Marshal(value) + if err != nil { + t.Fatalf("marshal %s: %v", label, err) + } + if strings.Contains(string(raw), forbidden) { + t.Fatalf("%s leaked %q: %s", label, forbidden, string(raw)) + } +} + func TestAdminExposesConfigSchemaQuotaAndDetectedModel(t *testing.T) { driver := &runtimeAdminDriver{} sdk := agentadaptor.New( diff --git a/util.go b/util.go index 44d5f4a..b88b9ae 100644 --- a/util.go +++ b/util.go @@ -31,6 +31,17 @@ func cloneStringMap(values map[string]string) map[string]string { return out } +func cloneEnvBindings(values []EnvBinding) []EnvBinding { + if len(values) == 0 { + return nil + } + out := make([]EnvBinding, 0, len(values)) + for _, value := range values { + out = append(out, EnvBinding{Name: value.Name, Value: value.Value}) + } + return out +} + func cloneAnyMap(values map[string]any) map[string]any { if len(values) == 0 { return nil @@ -99,6 +110,7 @@ func cloneRuntimePayload(payload RuntimePayload) RuntimePayload { return RuntimePayload{ Requested: cloneRuntimeServiceSpecs(payload.Requested), Ensured: cloneRuntimeServiceRefs(payload.Ensured), + SecretEnv: cloneEnvBindings(payload.SecretEnv), Fingerprint: payload.Fingerprint, } } diff --git a/workspace_skill_types.go b/workspace_skill_types.go index c465836..cc29e6c 100644 --- a/workspace_skill_types.go +++ b/workspace_skill_types.go @@ -264,6 +264,10 @@ type RuntimeServiceRef struct { OwnerAgentID string Health RuntimeServiceHealth Metadata map[string]string + // SecretEnv is a subprocess-only channel for runtime-issued secrets, such + // as per-run MCP bearer tokens. The SDK strips these bindings from public + // runtime refs/reports and injects them only into adapter process env. + SecretEnv []EnvBinding } // RuntimePayload is the runtime-service equivalent of ResolvedSkills. @@ -274,6 +278,7 @@ type RuntimeServiceRef struct { type RuntimePayload struct { Requested []RuntimeServiceSpec Ensured []RuntimeServiceRef + SecretEnv []EnvBinding Fingerprint string }