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..cb41e6d 100644 --- a/docs/streaming.md +++ b/docs/streaming.md @@ -145,11 +145,46 @@ 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()). 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`. +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 +204,7 @@ handle2, _ := sdk.Start(ctx, "batch", 覆盖顺序：per-call `WithStreaming` / `WithoutStreaming` > per-binding `WithDefaultStreaming` > 默认 off。 -### 4.2 背压 +### 5.2 背压 ```go sdk := agentadaptor.New( @@ -183,7 +218,7 @@ sdk := agentadaptor.New( **宿主务必及时 drain `StreamEvents()`**，否则 Block 模式可能挂起，Drop 模式会丢事件。 -## 5. 能力与降级 +## 6. 能力与降级 每个 streaming-aware adapter 通过 `StreamAwareDriver.StreamCapability()` 声明自己能提供什么： @@ -199,7 +234,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 +251,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/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/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..91b994d 100644 --- a/pkg/bridges/sse/handler.go +++ b/pkg/bridges/sse/handler.go @@ -38,12 +38,15 @@ import ( "io" "net/http" "strings" + "sync" "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" "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 +91,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. @@ -150,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.Protocol) + 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() } }) } @@ -235,25 +249,33 @@ 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, writeMu *sync.Mutex, 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 { case <-ctx.Done(): 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", protocol) + return fmt.Errorf("sse: unknown protocol %d", opts.Protocol) } } @@ -261,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() { @@ -291,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) + } } } } @@ -319,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 { @@ -327,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/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/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 new file mode 100644 index 0000000..f3f9b98 --- /dev/null +++ b/pkg/bridges/subagentstream/mux.go @@ -0,0 +1,325 @@ +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 +} + +const terminalFlushTimeout = 250 * time.Millisecond + +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 + } + if !sendAGUIWithCancelGrace(ctx, out, ev.AGUI) { + return + } + } + }() + return out +} + +func Wrap(ctx context.Context, handle agentadaptor.RunHandle, opts MuxOptions) <-chan Event { + out := make(chan Event, 64) + var seq atomic.Uint64 + + 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 out <- wrapped: + return true + case <-time.After(terminalFlushTimeout): + return false + } + } + 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 { + return sendSubagentEvent(ev, true) + } + active := newDelegationTracker() + drainSubagents := func(subagents <-chan a2adelegation.DelegationEvent) bool { + for subagents != nil { + select { + case ev, ok := <-subagents: + if !ok { + return true + } + active.Track(ev) + 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 + var cancelSubagents context.CancelFunc + if opts.Bus != nil { + 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(): + 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 + case ev, ok := <-parent: + if !ok { + parent = nil + if !drainSubagents(subagents) { + return + } + if !active.FlushSynthetic(a2adelegation.DelegationFailed, "failed", parentFinishedError(), sendSubagent) { + return + } + stopSubagents() + continue + } + if isTerminalAGUI(ev) { + if !drainSubagents(subagents) { + return + } + if !active.FlushSynthetic(a2adelegation.DelegationFailed, "failed", parentFinishedError(), sendSubagent) { + return + } + if !sendAGUI(ev) { + return + } + parent = nil + stopSubagents() + continue + } + if !sendAGUI(ev) { + return + } + case ev, ok := <-subagents: + if !ok { + subagents = nil + continue + } + active.Track(ev) + if !sendSubagent(ev) { + return + } + } + } + }() + 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() + _ = 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 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, + "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 + } + 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..df090ba --- /dev/null +++ b/pkg/bridges/subagentstream/mux_test.go @@ -0,0 +1,607 @@ +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 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{ + 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 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()) + 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 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) + 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(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 + readyOnce sync.Once + runID string + ctxDone <-chan struct{} + ready chan struct{} + events chan a2adelegation.DelegationEvent +} + +func newTrackingBus() *trackingBus { + return &trackingBus{ready: make(chan struct{}), events: make(chan a2adelegation.DelegationEvent, 16)} +} + +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() + 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() + 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 + 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..1ed2049 --- /dev/null +++ b/pkg/hosttools/a2adelegation/delegator.go @@ -0,0 +1,523 @@ +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 +} + +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) (A2AStream, 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"} + + 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) + 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} + 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 + } + + 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, + HistoryLength: req.HistoryLength, + 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 + lastTaskID := "" + cancelResult := func() (DelegationResult, error) { + if lastTaskID != "" { + d.cancelRemote(ctx, client, lastTaskID, send.Tenant, baseEvent) + } 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(ctx, client, lastTaskID, send.Tenant, baseEvent) + } + } + for { + select { + case <-ctx.Done(): + return cancelResult() + 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(): + _ = stream.Close() + return cancelResult() + case received := <-item: + event = received.event + err = received.err + } + if err == io.EOF { + 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, send.HistoryLength); ok { + lastTask = recovered + for _, ev := range mapper.taskEvents(recovered) { + d.publish(ev) + } + return d.finishTask(baseEvent, baseResult, recovered, spec.Policy, maxArtifacts) + } + cancelKnownTask() + } + 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.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, event.TaskID, send.Tenant, send.HistoryLength) + 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 + } + 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)} + } + return baseResult, terminalError(event.Status.State, spec.Policy) + } + } + if lastTask.ID != "" && executionFinalState(lastTask.Status.State) { + return d.finishTask(baseEvent, baseResult, lastTask, spec.Policy, maxArtifacts) + } + cancelKnownTask() + 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 +} + +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) + 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(ctx, 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, HistoryLength: send.HistoryLength}) + 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(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 + return baseResult, derr +} + +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, HistoryLength: historyLength}) + return task, err == nil && executionFinalState(task.Status.State) +} + +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" + return ev +} + +func (d *Delegator) cancelRemoteTask(ctx context.Context, client A2AClient, taskID, tenant string, base DelegationEvent) { + 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) { + if d.Bus != nil { + d.Bus.Publish(ev) + } +} + +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, + 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 { + if d != nil && d.NewID != nil { + return d.NewID() + } + 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) { + 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) + mapper := newEventMapper(baseEvent) + d.publish(mapper.terminalForState(task.ID, task.ContextID, task.Status.State, task.Raw)) + 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..e6ca198 --- /dev/null +++ b/pkg/hosttools/a2adelegation/delegator_test.go @@ -0,0 +1,1085 @@ +package a2adelegation + +import ( + "context" + "encoding/json" + "errors" + "io" + "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 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() + 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.terminalForState("task-1", "ctx-1", clienta2a.TaskStateInputRequired, nil) + if terminal.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" } + historyLength := 7 + + 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) + } + 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) + } + 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) + } +} + +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 + wantTerminal DelegationEventKind + }{ + {name: "default rejects", allow: false, wantErr: true, wantStatus: "failed", wantTerminal: DelegationFailed}, + {name: "policy allows", allow: true, wantErr: false, wantStatus: "input_required", wantTerminal: DelegationInputRequired}, + } { + 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, 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 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) + } + 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) + } + }) + } +} + +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) + } + 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, 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.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) + } + 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) { + 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 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}} + 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 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}}} + 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(), 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) + } + replayed := drainBus(t, bus, "run-1", 3) + if replayed[len(replayed)-1].Kind != DelegationCancelled { + t.Fatalf("expected cancelled terminal, got %#v", replayed) + } +} + +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}} + 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 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}} + 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) + 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 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) + 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()) + 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 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 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}} + 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,"history_length":3}}}}` + 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) + } + 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) { + 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) + } +} + +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 + + 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 + cancelHadDeadline bool +} + +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++ + f.lastSend = req + return f.sendTask, nil +} + +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, req clienta2a.GetTaskRequest) (clienta2a.Task, error) { + f.getCalls++ + f.lastGet = req + 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(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 +} + +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 new file mode 100644 index 0000000..9dc9866 --- /dev/null +++ b/pkg/hosttools/a2adelegation/events.go @@ -0,0 +1,140 @@ +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] { + 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 + } + 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) + 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() + 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: + 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..5a65410 --- /dev/null +++ b/pkg/hosttools/a2adelegation/mapping.go @@ -0,0 +1,308 @@ +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: + if event.Message != nil { + out = append(out, m.messageEvents(*event.Message)...) + } + } + 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)) + } + 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..a1fac99 --- /dev/null +++ b/pkg/hosttools/a2adelegation/mcpserver.go @@ -0,0 +1,342 @@ +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"` + 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"` + 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 { + 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}, + "history_length": 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["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 +} + +type MCPServerOptions struct { + RunID string + 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 { + 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 { + 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 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": + 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) + } + 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: 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.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) { + 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..bb907cf --- /dev/null +++ b/pkg/hosttools/a2adelegation/types.go @@ -0,0 +1,239 @@ +// 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 + HistoryLength *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.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 d3f9eee..aaa399d 100644 --- a/runtime_admin_test.go +++ b/runtime_admin_test.go @@ -2,13 +2,16 @@ package agentadaptor_test import ( "context" + "encoding/json" + "strings" "testing" agentadaptor "github.com/agent-dance/agent-adaptor" ) type runtimeAdminDriver struct { - lastReq agentadaptor.DriverRunRequest + lastReq agentadaptor.DriverRunRequest + mcpCapability agentadaptor.MCPCapability } func (d *runtimeAdminDriver) Descriptor() agentadaptor.DriverDescriptor { @@ -21,6 +24,7 @@ func (d *runtimeAdminDriver) Descriptor() agentadaptor.DriverDescriptor { }, }, Runtime: agentadaptor.RuntimeCapability{ReportsServices: true}, + MCP: d.mcpCapability, } } @@ -189,6 +193,252 @@ func TestRuntimeServicesFlowThroughSingleExecutionPath(t *testing.T) { } } +func TestRuntimeServiceMetadataInjectsMCPIntoDriverAndProfile(t *testing.T) { + const secret = "sk-runtime-secret" + + 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.visibility": "public-metadata", + }, + SecretEnv: []agentadaptor.EnvBinding{{Name: "DELEGATION_TOKEN", Value: 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), + ) + + result, err := sdk.Run(context.Background(), "use delegation") + if 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 !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) + } + 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 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}, + } + 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 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 }