diff --git a/agent-schema.json b/agent-schema.json index d73d6591e..44a6bbc8c 100644 --- a/agent-schema.json +++ b/agent-schema.json @@ -798,6 +798,15 @@ "description": "Enable strict schema adherence (OpenAI only). When true, all properties must be in required array.", "default": false }, + "mode": { + "type": "string", + "enum": [ + "native", + "tool" + ], + "description": "How the schema is enforced. 'native' (default) passes it to the provider's native structured-output support. 'tool' exposes an internal tool the model must call with the final JSON; the runtime validates the arguments against the schema. Works with any tool-calling provider.", + "default": "native" + }, "schema": { "type": "object", "description": "JSON Schema object defining the structure of the response. Must include type, properties, and required fields.", diff --git a/docs/configuration/structured-output/index.md b/docs/configuration/structured-output/index.md index ee370e7d6..9dc4c9e30 100644 --- a/docs/configuration/structured-output/index.md +++ b/docs/configuration/structured-output/index.md @@ -57,12 +57,45 @@ agents: ## Properties -| Property | Type | Required | Description | -| ------------- | ------- | -------- | --------------------------------------------------- | -| `name` | string | ✓ | Name identifier for the output schema | -| `description` | string | ✗ | Description of what the output represents | -| `strict` | boolean | ✗ | Enforce strict schema validation (default: `false`) | -| `schema` | object | ✓ | JSON Schema defining the output structure | +| Property | Type | Required | Description | +| ------------- | ------- | -------- | ----------------------------------------------------------------- | +| `name` | string | ✓ | Name identifier for the output schema | +| `description` | string | ✗ | Description of what the output represents | +| `strict` | boolean | ✗ | Enforce strict schema adherence — `native` mode only (default: `false`) | +| `schema` | object | ✓ | JSON Schema defining the output structure | +| `mode` | string | ✗ | Enforcement mode: `native` (default) or `tool` (see [Modes](#modes)) | + +## Modes + +### `native` (default) + +The schema is passed to the provider's native structured-output support (OpenAI JSON mode, Gemini JSON mode, ...). Omitting `mode` keeps this behavior. + +### `tool` + +```yaml +structured_output: + mode: tool + name: analysis_result + schema: + type: object + properties: + summary: + type: string + required: ["summary"] +``` + +In tool mode nothing is sent to the provider's native structured-output API. Instead, the runtime exposes an internal tool named `__structured_output__` whose parameters are exactly the configured schema. The model works normally — including calling other tools — and delivers its final answer by calling that tool, alone, as the only tool call of its response. The runtime validates the arguments against the schema: + +- A valid call ends the turn; the validated (compacted) JSON becomes the final assistant message. +- Invalid JSON gets a detailed tool error so the model can correct itself and retry. +- If the model answers in plain text instead, the runtime injects a transient system reminder and retries (at most 2 reminders), then fails with a `structured_output_failed` error. + +Tool-mode validation applies the full JSON Schema, including `additionalProperties` — unexpected fields are rejected when the schema forbids them. External `$ref` references (`http(s)://`, `file://`, cross-document) are rejected when the schema is compiled; only same-document references starting with `#` (e.g. `#/definitions/item`) are allowed. The `strict` flag has no effect in tool mode. + +Fork-mode skills (`context: fork`) run as exempt sub-sessions: the skill produces its own plain-text answer for the calling agent and is not required to call the output tool. The parent agent still delivers its final answer through the tool. + +Use tool mode when the model must combine tool use with a schema-constrained final answer, or when the provider has no native structured-output support. See [`examples/structured-output-tool-mode.yaml`](https://github.com/docker/docker-agent/blob/main/examples/structured-output-tool-mode.yaml). ## Schema Format @@ -119,7 +152,7 @@ schema: ## Strict Mode -When `strict: true`, the model is constrained to only produce output that exactly matches the schema. This provides stronger guarantees but may limit the model's flexibility. +`strict` only applies to `native` mode: it is passed to the provider's structured-output API. Tool mode ignores it and always validates against the full schema instead. When `strict: true`, the model is constrained to only produce output that exactly matches the schema. This provides stronger guarantees but may limit the model's flexibility. - **`strict: false` (default)** — model aims to match the schema but may include additional fields or slight variations. - **`strict: true`** — model output is constrained to exactly match the schema. Stronger guarantees. @@ -214,4 +247,4 @@ agents: > [!WARNING] > **Tool Limitations** > -> When using structured output, the agent typically cannot use tools since its response format is constrained to the schema. Design your agent workflow accordingly — structured output agents work best for single-turn analysis or extraction tasks. +> When using native structured output, the agent typically cannot use tools since its response format is constrained to the schema. Design your agent workflow accordingly — native structured output agents work best for single-turn analysis or extraction tasks. Use `mode: tool` when the agent needs to call tools before producing its schema-constrained final answer. diff --git a/examples/structured-output-tool-mode.yaml b/examples/structured-output-tool-mode.yaml new file mode 100644 index 000000000..6b68a625b --- /dev/null +++ b/examples/structured-output-tool-mode.yaml @@ -0,0 +1,39 @@ +#!/usr/bin/env docker agent run + +# Tool-mode structured output: instead of asking the provider for native +# structured output, the runtime exposes an internal tool whose parameters +# are exactly the configured schema. The model keeps using its regular tools +# and ends the conversation by calling that tool with the final JSON, which +# the runtime validates against the schema. Works with any tool-calling +# model, even without native structured-output support. + +agents: + root: + model: openai/gpt-4o + description: Agent that inspects files and reports findings as JSON + toolsets: + - type: filesystem + instruction: | + Inspect the files the user asks about, then report your findings. + Deliver the final report through the structured output tool. + structured_output: + mode: tool + name: file_report + description: Findings about the inspected files + schema: + type: object + properties: + files: + type: array + items: + type: object + properties: + path: + type: string + summary: + type: string + required: ["path", "summary"] + overall_summary: + type: string + required: ["files", "overall_summary"] + additionalProperties: false diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index e787bddb1..b07a213e7 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -16,6 +16,7 @@ import ( "github.com/docker/docker-agent/pkg/config/types" "github.com/docker/docker-agent/pkg/model/provider" "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/builtin/structuredoutput" ) // Agent represents an AI agent @@ -54,6 +55,12 @@ type Agent struct { harness *latest.HarnessConfig hooks *latest.HooksConfig cache *cache.Cache + structuredOutput *latest.StructuredOutput + // structuredOutputTool lazily compiles the tool-mode output tool once + // per agent (sync.OnceValues over structuredoutput.New); nil when the + // agent has no tool-mode structured output. WithStructuredOutput + // replaces the closure, which resets the cache. + structuredOutputTool func() (*structuredoutput.OutputTool, error) // warningsMu guards pendingWarnings. AddToolWarning and DrainWarnings // may be called concurrently from the runtime loop, the MCP server, @@ -392,6 +399,25 @@ func (a *Agent) Hooks() *latest.HooksConfig { return a.hooks } +// StructuredOutput returns the agent's structured-output configuration, or +// nil when none was configured. The runtime uses it to enforce tool-mode +// structured output; native mode is handled by the providers themselves. +func (a *Agent) StructuredOutput() *latest.StructuredOutput { + return a.structuredOutput +} + +// StructuredOutputTool returns the compiled tool-mode structured-output +// tool, building it on first use and caching both the tool and the error +// (schema compilation is not free, and the runtime asks on every turn). +// It returns (nil, nil) when the agent has no tool-mode structured output +// configured; native mode never compiles anything here. +func (a *Agent) StructuredOutputTool() (*structuredoutput.OutputTool, error) { + if a.structuredOutputTool == nil { + return nil, nil + } + return a.structuredOutputTool() +} + // Cache returns the response cache configured for this agent, or nil when // caching is disabled. func (a *Agent) Cache() *cache.Cache { diff --git a/pkg/agent/opts.go b/pkg/agent/opts.go index bd7464620..51919c58b 100644 --- a/pkg/agent/opts.go +++ b/pkg/agent/opts.go @@ -1,6 +1,7 @@ package agent import ( + "sync" "time" "github.com/docker/docker-agent/pkg/cache" @@ -8,6 +9,7 @@ import ( "github.com/docker/docker-agent/pkg/config/types" "github.com/docker/docker-agent/pkg/model/provider" "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/builtin/structuredoutput" ) type Opt func(a *Agent) @@ -259,6 +261,22 @@ func WithHooks(hooks *latest.HooksConfig) Opt { } } +// WithStructuredOutput keeps the agent's original structured-output +// configuration so the runtime can enforce tool-mode structured output. +// In tool mode it also (re)arms the lazy compile-once cache behind +// [Agent.StructuredOutputTool], so a reconfigured schema is recompiled. +func WithStructuredOutput(structuredOutput *latest.StructuredOutput) Opt { + return func(a *Agent) { + a.structuredOutput = structuredOutput + a.structuredOutputTool = nil + if structuredOutput.ToolMode() { + a.structuredOutputTool = sync.OnceValues(func() (*structuredoutput.OutputTool, error) { + return structuredoutput.New(structuredOutput) + }) + } + } +} + // WithCache attaches a response cache to the agent. Pass nil to disable. func WithCache(c *cache.Cache) Opt { return func(a *Agent) { diff --git a/pkg/agent/opts_test.go b/pkg/agent/opts_test.go index 6102f00a5..b0ae696f4 100644 --- a/pkg/agent/opts_test.go +++ b/pkg/agent/opts_test.go @@ -8,6 +8,7 @@ import ( "github.com/stretchr/testify/require" + "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/tools" ) @@ -96,3 +97,68 @@ func TestStartableToolSet_RestartAfterStop(t *testing.T) { require.True(t, ts.IsStarted()) require.Equal(t, int64(2), inner.starts.Load()) } + +// TestWithStructuredOutput verifies the original structured-output config is +// retained on the agent so the runtime can enforce tool mode. +func TestWithStructuredOutput(t *testing.T) { + t.Parallel() + + so := &latest.StructuredOutput{ + Name: "result", + Schema: map[string]any{"type": "object"}, + Mode: latest.StructuredOutputModeTool, + } + + a := New("root", "prompt", WithStructuredOutput(so)) + require.Same(t, so, a.StructuredOutput()) + + require.Nil(t, New("bare", "prompt").StructuredOutput()) +} + +// TestStructuredOutputTool_CompiledOnceAndResetByOpt pins the lazy cache +// contract: the tool-mode output tool is compiled once and reused across +// calls, WithStructuredOutput resets the cache, native mode (and no +// structured output) compiles nothing, and compile errors surface to the +// caller. +func TestStructuredOutputTool_CompiledOnceAndResetByOpt(t *testing.T) { + t.Parallel() + + toolMode := &latest.StructuredOutput{ + Name: "result", + Mode: latest.StructuredOutputModeTool, + Schema: map[string]any{"type": "object"}, + } + + a := New("root", "prompt", WithStructuredOutput(toolMode)) + first, err := a.StructuredOutputTool() + require.NoError(t, err) + require.NotNil(t, first) + second, err := a.StructuredOutputTool() + require.NoError(t, err) + require.Same(t, first, second, "repeated calls must return the cached tool") + + WithStructuredOutput(toolMode)(a) + third, err := a.StructuredOutputTool() + require.NoError(t, err) + require.NotSame(t, first, third, "WithStructuredOutput must reset the compile cache") + + native := New("native", "prompt", WithStructuredOutput(&latest.StructuredOutput{ + Name: "result", + Schema: map[string]any{"type": "object"}, + })) + tool, err := native.StructuredOutputTool() + require.NoError(t, err) + require.Nil(t, tool, "native mode must not compile the output tool") + + bare, err := New("bare", "prompt").StructuredOutputTool() + require.NoError(t, err) + require.Nil(t, bare) + + broken := New("broken", "prompt", WithStructuredOutput(&latest.StructuredOutput{ + Name: "broken", + Mode: latest.StructuredOutputModeTool, + Schema: map[string]any{"type": 42}, + })) + _, err = broken.StructuredOutputTool() + require.ErrorContains(t, err, "schema") +} diff --git a/pkg/config/latest/types.go b/pkg/config/latest/types.go index 9efc5cb4b..3332addf4 100644 --- a/pkg/config/latest/types.go +++ b/pkg/config/latest/types.go @@ -1995,6 +1995,17 @@ func (t TaskBudget) MarshalJSON() ([]byte, error) { return json.Marshal(map[string]any{"type": t.Type, "total": t.Total}) } +// Structured-output enforcement modes. See [StructuredOutput.Mode]. +const ( + // StructuredOutputModeNative passes the schema to the provider's + // native structured-output support (the default). + StructuredOutputModeNative = "native" + // StructuredOutputModeTool withholds the schema from providers and + // instead exposes an internal tool the model must call with the final + // JSON; the runtime validates the arguments against the schema. + StructuredOutputModeTool = "tool" +) + // StructuredOutput defines a JSON schema for structured output type StructuredOutput struct { // Name is the name of the response format @@ -2005,6 +2016,32 @@ type StructuredOutput struct { Schema map[string]any `json:"schema"` // Strict enables strict schema adherence (OpenAI only) Strict bool `json:"strict,omitempty"` + // Mode selects how the schema is enforced: "native" (or empty, the + // default) uses the provider's native structured-output support; + // "tool" makes the runtime expose an internal tool the model calls + // with the final JSON, validated against Schema. + Mode string `json:"mode,omitempty"` +} + +// ToolMode reports whether s opted into tool-based structured-output +// enforcement. Safe on a nil receiver so callers can chain it off +// optional config fields. +func (s *StructuredOutput) ToolMode() bool { + return s != nil && s.Mode == StructuredOutputModeTool +} + +// validate rejects unknown structured_output modes. nil is valid (no +// structured output configured). +func (s *StructuredOutput) validate() error { + if s == nil { + return nil + } + switch s.Mode { + case "", StructuredOutputModeNative, StructuredOutputModeTool: + return nil + default: + return fmt.Errorf("unsupported mode %q (must be one of: native, tool)", s.Mode) + } } // RAGToolConfig represents tool-specific configuration for a RAG source diff --git a/pkg/config/latest/validate.go b/pkg/config/latest/validate.go index 89647c767..168254519 100644 --- a/pkg/config/latest/validate.go +++ b/pkg/config/latest/validate.go @@ -78,6 +78,9 @@ func (t *Config) Validate() error { if err := agent.Safety.Validate(); err != nil { return fmt.Errorf("agents.%s.safety: %w", agent.Name, err) } + if err := agent.StructuredOutput.validate(); err != nil { + return fmt.Errorf("agents.%s.structured_output: %w", agent.Name, err) + } if err := validateCompactionThreshold(agent.CompactionThreshold); err != nil { return fmt.Errorf("agents.%s: %w", agent.Name, err) } @@ -218,6 +221,13 @@ func (a *AgentConfig) validateHarness() error { return errors.New("compaction_model cannot be used with a harness; the harness manages its own context compaction") } + // Harness agents also skip docker-agent's run loop, so the internal tool + // behind tool-mode structured output is never injected or enforced. + // Native mode stays accepted: the harness may support it itself. + if a.StructuredOutput.ToolMode() { + return errors.New("structured_output.mode 'tool' cannot be used with a harness; the harness does not run the tool loop that enforces it (use mode 'native')") + } + h := a.Harness switch h.Type { case "claude-code", "codex", "pi", "opencode": diff --git a/pkg/config/latest/validate_test.go b/pkg/config/latest/validate_test.go index 584bf53da..eddc8c7da 100644 --- a/pkg/config/latest/validate_test.go +++ b/pkg/config/latest/validate_test.go @@ -123,10 +123,11 @@ func TestAgentConfigValidateHarness(t *testing.T) { t.Parallel() tests := []struct { - name string - harness *HarnessConfig - compactionModel string - wantErr string + name string + harness *HarnessConfig + compactionModel string + structuredOutput *StructuredOutput + wantErr string }{ {name: "nil harness", harness: nil}, {name: "nil harness with compaction model", harness: nil, compactionModel: "fast"}, @@ -148,12 +149,16 @@ func TestAgentConfigValidateHarness(t *testing.T) { {name: "thinking on opencode", harness: &HarnessConfig{Type: "opencode", Thinking: true}}, {name: "thinking on wrong type", harness: &HarnessConfig{Type: "codex", Thinking: true}, wantErr: "harness.thinking can only be used with harness.type 'opencode'"}, {name: "compaction model on harness", harness: &HarnessConfig{Type: "claude-code"}, compactionModel: "fast", wantErr: "compaction_model cannot be used with a harness"}, + {name: "nil harness with tool-mode structured output", harness: nil, structuredOutput: &StructuredOutput{Mode: StructuredOutputModeTool}}, + {name: "native structured output on harness", harness: &HarnessConfig{Type: "claude-code"}, structuredOutput: &StructuredOutput{Mode: StructuredOutputModeNative}}, + {name: "default-mode structured output on harness", harness: &HarnessConfig{Type: "claude-code"}, structuredOutput: &StructuredOutput{Name: "out"}}, + {name: "tool-mode structured output on harness", harness: &HarnessConfig{Type: "claude-code"}, structuredOutput: &StructuredOutput{Mode: StructuredOutputModeTool}, wantErr: "structured_output.mode 'tool' cannot be used with a harness"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - agent := AgentConfig{Name: "root", Harness: tt.harness, CompactionModel: tt.compactionModel} + agent := AgentConfig{Name: "root", Harness: tt.harness, CompactionModel: tt.compactionModel, StructuredOutput: tt.structuredOutput} err := agent.validateHarness() if tt.wantErr != "" { require.ErrorContains(t, err, tt.wantErr) @@ -669,3 +674,68 @@ func TestSafetyModeValidate(t *testing.T) { } } } + +func TestStructuredOutputValidateMode(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + so *StructuredOutput + wantErr string + }{ + {name: "nil is valid", so: nil}, + {name: "mode absent", so: &StructuredOutput{Name: "x", Schema: map[string]any{"type": "object"}}}, + {name: "mode native", so: &StructuredOutput{Name: "x", Mode: StructuredOutputModeNative}}, + {name: "mode tool", so: &StructuredOutput{Name: "x", Mode: StructuredOutputModeTool}}, + {name: "unknown mode", so: &StructuredOutput{Name: "x", Mode: "json"}, wantErr: `unsupported mode "json" (must be one of: native, tool)`}, + {name: "case sensitive", so: &StructuredOutput{Name: "x", Mode: "Tool"}, wantErr: "unsupported mode"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := tt.so.validate() + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + }) + } +} + +// TestConfigValidateStructuredOutputMode pins that an unknown +// structured_output.mode is rejected at config level with the agent's name +// in the error, and the two documented modes (plus absent) pass. +func TestConfigValidateStructuredOutputMode(t *testing.T) { + t.Parallel() + + build := func(mode string) *Config { + return &Config{Agents: []AgentConfig{{ + Name: "root", + StructuredOutput: &StructuredOutput{ + Name: "result", + Schema: map[string]any{"type": "object"}, + Mode: mode, + }, + }}} + } + + require.NoError(t, build("").Validate()) + require.NoError(t, build(StructuredOutputModeNative).Validate()) + require.NoError(t, build(StructuredOutputModeTool).Validate()) + + err := build("strict").Validate() + require.ErrorContains(t, err, "agents.root.structured_output") + require.ErrorContains(t, err, `unsupported mode "strict"`) +} + +func TestStructuredOutputToolMode(t *testing.T) { + t.Parallel() + + var nilSO *StructuredOutput + require.False(t, nilSO.ToolMode()) + require.False(t, (&StructuredOutput{}).ToolMode()) + require.False(t, (&StructuredOutput{Mode: StructuredOutputModeNative}).ToolMode()) + require.True(t, (&StructuredOutput{Mode: StructuredOutputModeTool}).ToolMode()) +} diff --git a/pkg/model/provider/options/options.go b/pkg/model/provider/options/options.go index 465f3c299..5ca7d2222 100644 --- a/pkg/model/provider/options/options.go +++ b/pkg/model/provider/options/options.go @@ -110,8 +110,17 @@ func WithGateway(gateway string) Opt { } } +// WithStructuredOutput sets the native structured-output configuration +// forwarded to providers. Tool-mode structured output +// (structured_output.mode: tool) is enforced by the runtime through an +// internal tool, so it is deliberately dropped here: no provider must +// constrain the response natively in that mode. func WithStructuredOutput(structuredOutput *latest.StructuredOutput) Opt { return func(cfg *ModelOptions) { + if structuredOutput.ToolMode() { + cfg.structuredOutput = nil + return + } cfg.structuredOutput = structuredOutput } } diff --git a/pkg/model/provider/options/options_test.go b/pkg/model/provider/options/options_test.go index 003bdf95a..d591fdfc4 100644 --- a/pkg/model/provider/options/options_test.go +++ b/pkg/model/provider/options/options_test.go @@ -6,6 +6,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/config/latest" ) // sentinelTransport is a minimal http.RoundTripper used only for identity checks. @@ -41,6 +43,48 @@ func TestTransportWrapper_NilByDefault(t *testing.T) { assert.Nil(t, opts.TransportWrapper()) } +func TestWithStructuredOutput_ModeHandling(t *testing.T) { + t.Parallel() + + schema := map[string]any{"type": "object"} + + tests := []struct { + name string + so *latest.StructuredOutput + wantNativ bool + }{ + {name: "nil clears", so: nil, wantNativ: false}, + {name: "mode absent forwarded", so: &latest.StructuredOutput{Name: "x", Schema: schema}, wantNativ: true}, + {name: "mode native forwarded", so: &latest.StructuredOutput{Name: "x", Schema: schema, Mode: latest.StructuredOutputModeNative}, wantNativ: true}, + {name: "mode tool suppressed", so: &latest.StructuredOutput{Name: "x", Schema: schema, Mode: latest.StructuredOutputModeTool}, wantNativ: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + opts := Apply(WithStructuredOutput(tt.so)) + if tt.wantNativ { + require.Same(t, tt.so, opts.StructuredOutput()) + } else { + assert.Nil(t, opts.StructuredOutput(), "providers must not receive native structured output") + } + }) + } +} + +// TestWithStructuredOutput_ToolModeOverridesEarlier proves a later tool-mode +// opt clears a previously set native structured output, matching the +// "later opts override earlier ones" contract. +func TestWithStructuredOutput_ToolModeOverridesEarlier(t *testing.T) { + t.Parallel() + + native := &latest.StructuredOutput{Name: "x", Schema: map[string]any{"type": "object"}} + toolMode := &latest.StructuredOutput{Name: "x", Schema: map[string]any{"type": "object"}, Mode: latest.StructuredOutputModeTool} + + opts := Apply(WithStructuredOutput(native), WithStructuredOutput(toolMode)) + assert.Nil(t, opts.StructuredOutput()) +} + func TestFromModelOptions_RoundTripsTransportWrapper(t *testing.T) { t.Parallel() var wrapperInvoked bool diff --git a/pkg/runtime/agent_delegation.go b/pkg/runtime/agent_delegation.go index 265eef498..92dd01782 100644 --- a/pkg/runtime/agent_delegation.go +++ b/pkg/runtime/agent_delegation.go @@ -161,6 +161,11 @@ type SubSessionConfig struct { // top of the agent's own toolsets. Used by fork-mode skills that declare // assistive toolsets. ExtraToolSets []tools.ToolSet + // DisableStructuredOutput exempts the child session from tool-mode + // structured-output enforcement. Set only by fork-mode skills, whose + // answer is free-form text for the calling agent; ordinary delegations + // and background agents stay enforced. + DisableStructuredOutput bool // DelegationLineage is the chain of agents that delegated to produce this // child session (parent lineage plus the delegating caller), as computed // by validateDelegation. Set only for true delegation edges @@ -241,6 +246,7 @@ func newSubSession(parent *session.Session, cfg SubSessionConfig, childAgent *ag session.WithSafetyPolicy(cfg.SafetyPolicy), session.WithNonInteractive(cfg.NonInteractive), session.WithSendUserMessage(false), + session.WithStructuredOutputDisabled(cfg.DisableStructuredOutput), session.WithParentID(parent.ID), session.WithAttachedFiles(attachedFiles), session.WithAttributes(parent.AttributesSnapshot()), diff --git a/pkg/runtime/agent_delegation_test.go b/pkg/runtime/agent_delegation_test.go index 6bfe25ec6..399ca7313 100644 --- a/pkg/runtime/agent_delegation_test.go +++ b/pkg/runtime/agent_delegation_test.go @@ -172,6 +172,19 @@ func TestNewSubSession(t *testing.T) { // We can verify the user message is still the default. assert.Equal(t, "Please proceed.", s.GetLastUserMessageContent()) }) + + t.Run("disable structured output", func(t *testing.T) { + cfg := SubSessionConfig{ + Task: "run a fork skill", + AgentName: "worker", + Title: "Skill task", + DisableStructuredOutput: true, + } + + s := newSubSession(parent, cfg, childAgent) + + assert.True(t, s.DisableStructuredOutput) + }) } func TestSubSessionConfig_DefaultValues(t *testing.T) { @@ -191,6 +204,7 @@ func TestSubSessionConfig_DefaultValues(t *testing.T) { assert.False(t, s.ToolsApproved) assert.False(t, s.SendUserMessage) + assert.False(t, s.DisableStructuredOutput) assert.Empty(t, s.AgentName) } diff --git a/pkg/runtime/cache.go b/pkg/runtime/cache.go index 0b2a19dba..e26d5c853 100644 --- a/pkg/runtime/cache.go +++ b/pkg/runtime/cache.go @@ -77,6 +77,31 @@ func (r *LocalRuntime) tryReplayCachedResponse( return false } + // Tool-mode structured output only ever accepts schema-valid JSON as the + // final answer, so a cached entry must pass the same validation before + // being replayed. A stale entry (e.g. plain text stored before the agent + // switched to tool mode) is treated as a miss and the model is called; + // a valid one is replayed in its canonical compacted form. Sessions that + // disable structured output (fork-mode skill children) skip this and + // replay any cached text as-is. + if structuredOutputEnabled(sess, a) { + ot, otErr := a.StructuredOutputTool() + if otErr != nil { + // Uncompilable schema: fall through so getTools surfaces the + // configuration error on the regular model path. + return false + } + if ot != nil { + validated, err := ot.Validate(cached) + if err != nil { + slog.DebugContext(ctx, "Response cache hit is not valid structured output; treating as miss", + "agent", a.Name(), "session_id", sess.ID, "error", err) + return false + } + cached = validated + } + } + slog.DebugContext(ctx, "Response cache hit; replaying cached answer", "agent", a.Name(), "session_id", sess.ID) modelID := a.Model(ctx).ID().String() diff --git a/pkg/runtime/cache_test.go b/pkg/runtime/cache_test.go index 45a66edc8..41d91c77e 100644 --- a/pkg/runtime/cache_test.go +++ b/pkg/runtime/cache_test.go @@ -1,6 +1,8 @@ package runtime import ( + "bytes" + "encoding/json" "strings" "testing" @@ -15,13 +17,14 @@ import ( "github.com/docker/docker-agent/pkg/team" ) -func runWithCache(t *testing.T, c *cache.Cache, prov *messageRecordingProvider, sess *session.Session) []Event { +func runWithCache(t *testing.T, c *cache.Cache, prov *messageRecordingProvider, sess *session.Session, extraOpts ...agent.Opt) []Event { t.Helper() - root := agent.New("root", "You are a test agent", + opts := append([]agent.Opt{ agent.WithModel(prov), agent.WithCache(c), - ) + }, extraOpts...) + root := agent.New("root", "You are a test agent", opts...) tm := team.New(team.WithAgents(root)) rt, err := NewLocalRuntime(t.Context(), tm, WithSessionCompaction(false), WithModelStore(mockModelStore{})) @@ -241,6 +244,62 @@ func TestCache_EmptyResponseIsNotCached(t *testing.T) { assert.False(t, found, "empty assistant response must not appear in the cache") } +// TestCache_ToolModeStaleHitFallsThroughToModel verifies that with tool-mode +// structured output a cached plain-text answer (e.g. stored before the agent +// switched to tool mode) fails revalidation and is treated as a miss: the +// model is called, the valid output-tool call becomes the final answer, and +// the stop hook replaces the stale entry with the validated JSON. +func TestCache_ToolModeStaleHitFallsThroughToModel(t *testing.T) { + t.Parallel() + + c, err := cache.New(cache.Config{Enabled: true}) + require.NoError(t, err) + c.Store("Answer me", "a plain text answer from before tool mode") + + prov := &messageRecordingProvider{ + id: "test/mock-model", + streams: []*mockStream{outputCallStream("call_1", `{"answer":"fresh"}`)}, + } + sess := session.New(session.WithUserMessage("Answer me")) + events := runWithCache(t, c, prov, sess, agent.WithStructuredOutput(toolModeStructuredOutput())) + + require.Empty(t, errorEvents(events)) + prov.mu.Lock() + require.Len(t, prov.recordedMessages, 1, "a stale cache hit must fall through to the model") + prov.mu.Unlock() + assert.JSONEq(t, `{"answer":"fresh"}`, sess.GetLastAssistantMessageContent()) + + stored, ok := c.Lookup("Answer me") + require.True(t, ok) + assert.JSONEq(t, `{"answer":"fresh"}`, stored, "the validated JSON must replace the stale entry") +} + +// TestCache_ToolModeValidHitReplaysWithoutModel verifies that a cached +// schema-valid answer passes tool-mode revalidation and is replayed in its +// canonical compacted form without calling the model. +func TestCache_ToolModeValidHitReplaysWithoutModel(t *testing.T) { + t.Parallel() + + c, err := cache.New(cache.Config{Enabled: true}) + require.NoError(t, err) + c.Store("Answer me", "{\n \"answer\": \"cached\"\n}") + + prov := &messageRecordingProvider{id: "test/mock-model"} + sess := session.New(session.WithUserMessage("Answer me")) + events := runWithCache(t, c, prov, sess, agent.WithStructuredOutput(toolModeStructuredOutput())) + + require.Empty(t, errorEvents(events)) + prov.mu.Lock() + assert.Empty(t, prov.recordedMessages, "a valid structured hit must not call the model") + prov.mu.Unlock() + last := sess.GetLastAssistantMessageContent() + assert.JSONEq(t, `{"answer":"cached"}`, last) + // Canonicalization contract: the replayed content is the compacted form. + var compacted bytes.Buffer + require.NoError(t, json.Compact(&compacted, []byte(last))) + assert.Equal(t, compacted.String(), last, "the replayed content must be canonical compact JSON") +} + func hasAgentChoice(events []Event, content string) bool { for _, ev := range events { if ac, ok := ev.(*AgentChoiceEvent); ok && strings.Contains(ac.Content, content) { diff --git a/pkg/runtime/event.go b/pkg/runtime/event.go index e2da85565..56cc24f95 100644 --- a/pkg/runtime/event.go +++ b/pkg/runtime/event.go @@ -232,14 +232,15 @@ func AgentChoiceReasoning(agentName, sessionID, content string) Event { // messages for each shape (token-count overflow, wire-level body cap, // media-size rejection) instead of one generic "context window exceeded". const ( - ErrorCodeModelError = "model_error" - ErrorCodeRateLimited = "rate_limited" - ErrorCodeContextExceeded = "context_exceeded" // OverflowKindTokens - ErrorCodeRequestTooLarge = "request_too_large" // OverflowKindWire - ErrorCodeMediaTooLarge = "media_too_large" // OverflowKindMedia - ErrorCodeToolFailed = "tool_failed" - ErrorCodeHookBlocked = "hook_blocked" - ErrorCodeLoopDetected = "loop_detected" + ErrorCodeModelError = "model_error" + ErrorCodeRateLimited = "rate_limited" + ErrorCodeContextExceeded = "context_exceeded" // OverflowKindTokens + ErrorCodeRequestTooLarge = "request_too_large" // OverflowKindWire + ErrorCodeMediaTooLarge = "media_too_large" // OverflowKindMedia + ErrorCodeToolFailed = "tool_failed" + ErrorCodeHookBlocked = "hook_blocked" + ErrorCodeLoopDetected = "loop_detected" + ErrorCodeStructuredOutputFailed = "structured_output_failed" ) type ErrorEvent struct { diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 3a1ea428b..26ba8e51b 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -372,13 +372,11 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, r.emitAgentWarnings(a, sink) r.configureToolsetHandlers(a, sink) - agentTools, err := r.getTools(ctx, a, sessionSpan, sink, true) + agentTools, err := r.getTools(ctx, sess, a, sessionSpan, sink, true) if err != nil { sink.Emit(ErrorWithCodeForSession(sess.ID, ErrorCodeToolFailed, fmt.Sprintf("failed to get tools: %v", err))) return } - agentTools = filterExcludedTools(agentTools, sess.ExcludedTools) - agentTools = r.skillSubSessionTools(ctx, sess, a, agentTools, sink) // Record the catalogue size on the session span — answers "how // many tools could this turn actually use?" without having to @@ -473,19 +471,18 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, if a.Name() != ls.prevAgentName { ls.toolModelOverride = "" ls.prevTurnMadeToolCalls = false + ls.structuredOutputReminders = 0 ls.prevAgentName = a.Name() } r.emitAgentWarnings(a, sink) r.configureToolsetHandlers(a, sink) - agentTools, err := r.getTools(ctx, a, sessionSpan, sink, true) + agentTools, err := r.getTools(ctx, sess, a, sessionSpan, sink, true) if err != nil { sink.Emit(ErrorWithCodeForSession(sess.ID, ErrorCodeToolFailed, fmt.Sprintf("failed to get tools: %v", err))) return } - agentTools = filterExcludedTools(agentTools, sess.ExcludedTools) - agentTools = r.skillSubSessionTools(ctx, sess, a, agentTools, sink) // Emit updated tool count. After a ToolListChanged MCP notification // the cache is invalidated, so getTools above re-fetches from the @@ -635,6 +632,11 @@ type loopState struct { sessionStartSources []session.InstructionSource userPromptMsgs []chat.Message exitReason string + // structuredOutputReminders counts the transient reminders injected + // after a tool-mode turn stopped without calling the internal output + // tool. Bounded by maxStructuredOutputReminders; reset on agent switch + // so it never carries across agents. + structuredOutputReminders int // prevTurnMadeToolCalls reports whether the immediately preceding // turn emitted tool calls. Used to classify an empty trailing turn: // a clean stop right after tool work is benign (the model already @@ -758,9 +760,13 @@ func (r *LocalRuntime) runTurn( // against what the model already knows. Changes extend the conversation; // they never rewrite the frozen instruction prefix. turnStartMsgs := r.executeTurnStartHooks(ctx, sess, a, events) - legacyExtras := slices.Concat(ls.sessionStartLegacyMsgs, ls.userPromptMsgs, turnStartMsgs.legacyMessages()) - messages := r.messagesWithDynamicContext(ctx, sess, a, - instructionSources(ls.sessionStartMsgs, ls.userPromptMsgs, turnStartMsgs, ls.sessionStartSources...), legacyExtras) + // Pending tool-mode structured-output reminder rides with the transient + // system extras: threaded per call, never persisted as a user message. + reminderMsgs := ls.structuredOutputReminderMessages() + legacyExtras := slices.Concat(ls.sessionStartLegacyMsgs, ls.userPromptMsgs, turnStartMsgs.legacyMessages(), reminderMsgs) + sources := instructionSources(ls.sessionStartMsgs, ls.userPromptMsgs, turnStartMsgs, ls.sessionStartSources...) + sources = append(sources, instructionSource("runtime/structured-output", "structured-output reminder", reminderMsgs)) + messages := r.messagesWithDynamicContext(ctx, sess, a, sources, legacyExtras) slog.DebugContext(ctx, "Retrieved messages for processing", "agent", a.Name(), "message_count", len(messages)) // before_llm_call hooks fire just before the model is invoked. @@ -889,7 +895,13 @@ func (r *LocalRuntime) runTurn( // measure how much content was added by tool results. messageCountBeforeTools := len(sess.OwnMessages()) - stopRun, stopMsg := r.processToolCalls(ctx, sess, res.Calls, agentTools, events) + // Intercept internal structured-output calls before dispatch: they bypass + // approval and the user's tool hooks, and an exclusive valid call + // finalizes the turn (res is rewritten so the stop path below runs on + // the validated JSON). + dispatchCalls, soFinalized := r.handleStructuredOutputCalls(ctx, sess, a, &res, agentTools, modelID.String(), events) + + stopRun, stopMsg := r.processToolCalls(ctx, sess, dispatchCalls, agentTools, events) // Re-probe toolsets after tool calls: an install/setup tool call may // have made a previously-unavailable LSP or MCP connectable. reprobe() @@ -973,6 +985,19 @@ func (r *LocalRuntime) runTurn( } if res.Stopped { + // Tool-mode structured output never accepts a plain-text stop as the + // final answer: remind the model (bounded) or fail with a coded error. + if structuredOutputEnabled(sess, a) { + if !soFinalized { + ctrl, reason := r.structuredOutputStop(ctx, sess, a, ls, events) + endReason = reason + return ctrl + } + // Accepted result: give a potential next turn (follow-up, + // steered continuation) a fresh reminder budget. + ls.structuredOutputReminders = 0 + } + slog.DebugContext(ctx, "Conversation stopped", "agent", a.Name()) r.executeStopHooks(ctx, sess, a, res.Content, events) @@ -1276,16 +1301,30 @@ func (r *LocalRuntime) compactIfNeeded( r.compactWithReason(ctx, sess, "", compactionReasonThreshold, events) } -// getTools executes tool retrieval with automatic OAuth handling. +// getTools executes tool retrieval with automatic OAuth handling and applies +// the session's tool-visibility rules: exclusion filters, the skill +// sub-session allow-list and extra toolsets, then the internal tool-mode +// structured-output tool. The internal tool is appended after every session +// filter so a skill allow-list can never strip active enforcement, and it is +// omitted entirely for sessions that disable structured output (fork-mode +// skill children). // emitLifecycleEvents controls whether MCPInitStarted/Finished are emitted; // pass false when calling from reprobe to avoid spurious TUI spinner flicker. -func (r *LocalRuntime) getTools(ctx context.Context, a *agent.Agent, sessionSpan trace.Span, events EventSink, emitLifecycleEvents bool) ([]tools.Tool, error) { +func (r *LocalRuntime) getTools(ctx context.Context, sess *session.Session, a *agent.Agent, sessionSpan trace.Span, events EventSink, emitLifecycleEvents bool) ([]tools.Tool, error) { if emitLifecycleEvents && len(a.ToolSets()) > 0 { events.Emit(MCPInitStarted(a.Name())) defer func() { events.Emit(MCPInitFinished(a.Name())) }() } agentTools, err := a.Tools(ctx) + if err == nil { + agentTools = filterExcludedTools(agentTools, sess.ExcludedTools) + agentTools = r.skillSubSessionTools(ctx, sess, a, agentTools, events) + // Tool-mode structured output rides on the same error path: a name + // collision with the reserved internal tool or an uncompilable schema + // must fail the turn loudly, not mask one of the two tools. + agentTools, err = appendStructuredOutputTool(agentTools, sess, a) + } if err != nil { slog.ErrorContext(ctx, "Failed to get agent tools", "agent", a.Name(), "error", err) sessionSpan.RecordError(err) @@ -1547,13 +1586,11 @@ func (r *LocalRuntime) reprobe( sessionSpan trace.Span, events EventSink, ) { - updated, err := r.getTools(ctx, a, sessionSpan, events, false) + updated, err := r.getTools(ctx, sess, a, sessionSpan, events, false) if err != nil { slog.WarnContext(ctx, "reprobe: getTools failed", "agent", a.Name(), "error", err) return } - updated = filterExcludedTools(updated, sess.ExcludedTools) - updated = r.skillSubSessionTools(ctx, sess, a, updated, events) // Emit any pending warnings that getTools just generated. r.emitAgentWarnings(a, events) diff --git a/pkg/runtime/runtime_test.go b/pkg/runtime/runtime_test.go index 815346c0d..dcb589e0e 100644 --- a/pkg/runtime/runtime_test.go +++ b/pkg/runtime/runtime_test.go @@ -1014,7 +1014,7 @@ func TestGetTools_WarningHandling(t *testing.T) { sessionSpan := trace.SpanFromContext(t.Context()) // First call - tools1, err := rt.getTools(t.Context(), root, sessionSpan, NewChannelSink(events), true) + tools1, err := rt.getTools(t.Context(), session.New(), root, sessionSpan, NewChannelSink(events), true) require.NoError(t, err) require.Len(t, tools1, tt.wantToolCount) diff --git a/pkg/runtime/skill_runner.go b/pkg/runtime/skill_runner.go index 4ddcbb925..3d0f34c38 100644 --- a/pkg/runtime/skill_runner.go +++ b/pkg/runtime/skill_runner.go @@ -130,6 +130,10 @@ func (r *LocalRuntime) RunSkillFork(ctx context.Context, sess *session.Session, ExcludedTools: []string{skills.ToolNameRunSkill}, AllowedTools: prepared.AllowedTools, ExtraToolSets: prepared.ToolSets, + // A fork skill's answer is free-form text consumed by the calling + // agent, not the agent's schema-constrained final output, so the + // child is exempt from tool-mode structured-output enforcement. + DisableStructuredOutput: true, }, }) } diff --git a/pkg/runtime/structured_output.go b/pkg/runtime/structured_output.go new file mode 100644 index 000000000..639cb9033 --- /dev/null +++ b/pkg/runtime/structured_output.go @@ -0,0 +1,347 @@ +package runtime + +import ( + "context" + "errors" + "fmt" + "log/slog" + "strings" + "time" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/telemetry/genai" + "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/builtin/structuredoutput" +) + +// maxStructuredOutputReminders bounds how many transient system reminders a +// tool-mode agent gets when it finishes a turn without calling the internal +// output tool. Once exhausted, the run fails with +// [ErrorCodeStructuredOutputFailed] instead of looping forever. +const maxStructuredOutputReminders = 2 + +// structuredOutputReminderText nudges the model back to the output tool +// after it produced a plain-text final answer in tool mode. Injected as a +// transient system message, never persisted to the session. +const structuredOutputReminderText = "Reminder: this conversation requires structured output. " + + "Do not answer in plain text. Call the " + structuredoutput.ToolName + + " tool with your final answer as arguments matching its JSON schema; a valid call ends the turn." + +// structuredOutputReminderMessages returns the pending transient reminder as +// a one-element system-message slice, or nil when no reminder is due. The +// counter is reset on agent switch by the run loop. +func (ls *loopState) structuredOutputReminderMessages() []chat.Message { + if ls.structuredOutputReminders == 0 { + return nil + } + return []chat.Message{{Role: chat.MessageRoleSystem, Content: structuredOutputReminderText}} +} + +// structuredOutputEnabled reports whether tool-mode structured output is +// enforced for this session: the agent opted into tool mode and the session +// did not opt out. Fork-mode skill sub-sessions opt out because their answer +// is free-form text for the calling agent, not the agent's final output. +func structuredOutputEnabled(sess *session.Session, a *agent.Agent) bool { + return a.StructuredOutput().ToolMode() && !sess.DisableStructuredOutput +} + +// appendStructuredOutputTool exposes the internal structured-output tool when +// tool-mode enforcement is active for the session, reusing the tool compiled +// once on the agent instead of recompiling the schema every turn. Called +// after the session's exclusion and skill allow-list filters so enforcement +// can never be filtered away, and skipped entirely for sessions that disable +// structured output. A real tool already using the reserved name is a +// configuration error: silently keeping either one would mask the other, so +// the turn fails with a clear message instead. +func appendStructuredOutputTool(agentTools []tools.Tool, sess *session.Session, a *agent.Agent) ([]tools.Tool, error) { + if !structuredOutputEnabled(sess, a) { + return agentTools, nil + } + ot, err := a.StructuredOutputTool() + if err != nil { + return nil, fmt.Errorf("invalid structured_output configuration: %w", err) + } + if ot == nil { + return agentTools, nil + } + for _, t := range agentTools { + if t.Name == structuredoutput.ToolName { + return nil, fmt.Errorf("tool name %q is reserved for tool-mode structured output; rename or remove the conflicting tool", structuredoutput.ToolName) + } + } + return append(agentTools, ot.Definition()), nil +} + +// handleStructuredOutputCalls intercepts calls to the internal structured- +// output tool before the batch reaches the dispatcher, so they never go +// through approval or the user's tool hooks. It returns the calls that +// should still be dispatched normally and whether the turn was finalized +// with a validated result. +// +// Semantics: +// - a single, exclusive output call is validated against the schema. On +// success the canonical JSON becomes a new final assistant message and +// res is rewritten (Stopped=true, Content=JSON) so the regular stop path +// (stop hooks, force_handoff, follow-ups) runs on the validated output. +// On validation failure a detailed error tool-response lets the model +// retry on the next iteration. +// - output calls in a mixed or multi-output batch are each rejected with +// an error tool-response (never terminal); the remaining calls follow +// normal dispatch. +// +// Every intercepted call gets a tool response recorded, so the history never +// contains an orphaned tool call, and the same runtime.tool.call span and +// tool-call metric a dispatched call would get; approval and user tool hooks +// stay bypassed. +func (r *LocalRuntime) handleStructuredOutputCalls( + ctx context.Context, + sess *session.Session, + a *agent.Agent, + res *streamResult, + agentTools []tools.Tool, + modelID string, + events EventSink, +) (dispatchCalls []tools.ToolCall, finalized bool) { + if !structuredOutputEnabled(sess, a) { + return res.Calls, false + } + + var outputCalls, otherCalls []tools.ToolCall + for _, tc := range res.Calls { + if tc.Function.Name == structuredoutput.ToolName { + outputCalls = append(outputCalls, tc) + } else { + otherCalls = append(otherCalls, tc) + } + } + if len(outputCalls) == 0 { + return res.Calls, false + } + + tool, ok := findToolByName(agentTools, structuredoutput.ToolName) + if !ok { + // The tool wasn't offered this turn (defensive: injection happens + // after every session filter, so this should not occur when + // enforcement is active); let the dispatcher reject the calls as + // unavailable rather than inventing a definition here. + return res.Calls, false + } + + if len(outputCalls) > 1 || len(otherCalls) > 0 { + // A terminal result must be the sole call of its batch. Reject each + // output call so the model retries exclusively; other tools proceed. + slog.WarnContext(ctx, "Structured output call rejected: not exclusive in its batch", + "agent", a.Name(), "session_id", sess.ID, + "output_calls", len(outputCalls), "other_calls", len(otherCalls)) + rejectMsg := fmt.Sprintf( + "Structured output rejected: %s must be the only tool call in the response. "+ + "Finish any other tool use first, then call %s alone with the final answer.", + structuredoutput.ToolName, structuredoutput.ToolName) + for _, tc := range outputCalls { + r.rejectStructuredOutputCall(ctx, sess, a, tc, tool, rejectMsg, events) + } + return otherCalls, false + } + + tc := outputCalls[0] + result := r.runStructuredOutputCall(ctx, sess, a, tc, tool, events) + if result.IsError { + slog.DebugContext(ctx, "Structured output call failed validation; letting the model retry", + "agent", a.Name(), "session_id", sess.ID, "error", result.Output) + return nil, false + } + + // Terminal: promote the validated JSON to a new final assistant message + // so GetLastAssistantMessageContent() and every consumer built on it + // (MCP, API, CLI) read the structured result. + addAgentMessage(sess, a, &chat.Message{ + Role: chat.MessageRoleAssistant, + Content: result.Output, + CreatedAt: r.now().Format(time.RFC3339), + Model: modelID, + }, events) + res.Content = result.Output + res.Stopped = true + return nil, true +} + +// runStructuredOutputCall executes one exclusive output call through the +// instrumented pipeline and records its ToolCall/ToolCallResponse pair. A +// handler error is folded into an IsError result the model can react to, +// mirroring the dispatcher's error translation. +func (r *LocalRuntime) runStructuredOutputCall( + ctx context.Context, + sess *session.Session, + a *agent.Agent, + tc tools.ToolCall, + tool tools.Tool, + events EventSink, +) *tools.ToolCallResult { + result := r.instrumentStructuredOutputCall(ctx, sess, a, tc, func(ctx context.Context) (*tools.ToolCallResult, error) { + return tool.Handler(ctx, tc, tools.NopRuntime{}) + }) + r.recordStructuredOutputResponse(sess, a, tc, tool, result, events) + return result +} + +// rejectStructuredOutputCall records a non-exclusive output call as rejected +// without running the tool handler. The rejection still goes through the +// instrumented pipeline so every intercepted call is observable, and each +// call gets its own result value so emitted events never share a pointer. +func (r *LocalRuntime) rejectStructuredOutputCall( + ctx context.Context, + sess *session.Session, + a *agent.Agent, + tc tools.ToolCall, + tool tools.Tool, + rejectMsg string, + events EventSink, +) { + result := r.instrumentStructuredOutputCall(ctx, sess, a, tc, func(context.Context) (*tools.ToolCallResult, error) { + return tools.ResultError(rejectMsg), nil + }) + r.recordStructuredOutputResponse(sess, a, tc, tool, result, events) +} + +// instrumentStructuredOutputCall wraps one intercepted output call with the +// observability the dispatcher gives regular tool calls: a runtime.tool.call +// span carrying the same gen_ai.* and legacy attributes, opt-in +// argument/result content capture, and exactly one RecordToolCall metric. +// It instruments only — approval and user tool hooks stay bypassed. +// +// Unlike the dispatcher, an IsError result counts as a failed call: for this +// internal tool it always means a schema violation or a batch rejection, so +// a synthetic error carries the rejection text to the metric and the span. +func (r *LocalRuntime) instrumentStructuredOutputCall( + ctx context.Context, + sess *session.Session, + a *agent.Agent, + tc tools.ToolCall, + exec func(ctx context.Context) (*tools.ToolCallResult, error), +) *tools.ToolCallResult { + attrs := []attribute.KeyValue{ + attribute.String(genai.AttrOperationName, genai.OperationExecuteTool), + attribute.String(genai.AttrToolName, tc.Function.Name), + attribute.String(genai.AttrToolType, "function"), + attribute.String(genai.AttrToolCallID, tc.ID), + attribute.String(genai.AttrAgentNameRuntime, a.Name()), + attribute.String(genai.AttrConversationID, sess.ID), + } + attrs = append(attrs, genai.LegacyToolAttributes( + tc.Function.Name, string(tc.Type), a.Name(), sess.ID, tc.ID, + )...) + ctx, span := r.startSpan(ctx, "runtime.tool.call", trace.WithAttributes(attrs...)) + defer span.End() + + if genai.IsContentCaptureEnabled() && tc.Function.Arguments != "" { + span.SetAttributes(attribute.String(genai.AttrToolCallArguments, tc.Function.Arguments)) + } + + start := r.now() + result, err := exec(ctx) + if err != nil { + result = tools.ResultError(fmt.Sprintf("Error calling tool: %v", err)) + } + callErr := err + if callErr == nil && result.IsError { + callErr = errors.New(result.Output) + } + r.telemetry.RecordToolCall(ctx, tc.Function.Name, sess.ID, a.Name(), r.now().Sub(start), callErr) + + switch { + case err != nil: + span.RecordError(err) + span.SetStatus(codes.Error, "structured output handler error") + case result.IsError: + span.RecordError(callErr) + span.SetStatus(codes.Error, "structured output rejected") + default: + span.SetStatus(codes.Ok, "structured output call completed") + } + + if genai.IsContentCaptureEnabled() && result.Output != "" { + span.SetAttributes(attribute.String(genai.AttrToolCallResult, result.Output)) + } + return result +} + +// recordStructuredOutputResponse emits the ToolCall/ToolCallResponse pair for +// an internal structured-output call and records the tool message in the +// session. It deliberately bypasses the dispatcher: the internal tool needs +// no approval and must not reach the user's tool hooks. +func (r *LocalRuntime) recordStructuredOutputResponse( + sess *session.Session, + a *agent.Agent, + tc tools.ToolCall, + tool tools.Tool, + result *tools.ToolCallResult, + events EventSink, +) { + events.Emit(ToolCall(tc, tool, a.Name())) + events.Emit(ToolCallResponse(tc.ID, tool, result, result.Output, a.Name())) + content := result.Output + if strings.TrimSpace(content) == "" { + content = "(no output)" + } + addAgentMessage(sess, a, &chat.Message{ + Role: chat.MessageRoleTool, + Content: content, + ToolCallID: tc.ID, + IsError: result.IsError, + CreatedAt: r.now().Format(time.RFC3339), + }, events) +} + +// structuredOutputStop decides what to do when a tool-mode turn stopped +// without a validated structured output: retry with a transient system +// reminder (at most maxStructuredOutputReminders times), then fail the run +// with a coded error. Returns turnContinue/turnExit and the matching +// turn-end reason. +func (r *LocalRuntime) structuredOutputStop( + ctx context.Context, + sess *session.Session, + a *agent.Agent, + ls *loopState, + events EventSink, +) (turnControl, string) { + if ls.structuredOutputReminders < maxStructuredOutputReminders { + ls.structuredOutputReminders++ + slog.InfoContext(ctx, "Model stopped without structured output; injecting transient reminder", + "agent", a.Name(), "session_id", sess.ID, + "reminder", ls.structuredOutputReminders, "max", maxStructuredOutputReminders) + return turnContinue, turnEndReasonContinue + } + errMsg := fmt.Sprintf( + "Agent terminated: the model did not deliver structured output via the %s tool after %d reminders.", + structuredoutput.ToolName, maxStructuredOutputReminders) + slog.WarnContext(ctx, "Structured output failed: reminders exhausted", + "agent", a.Name(), "session_id", sess.ID, "max", maxStructuredOutputReminders) + events.Emit(ErrorWithCodeForSession(sess.ID, ErrorCodeStructuredOutputFailed, errMsg)) + r.notifyError(ctx, a, sess.ID, errMsg) + // The non-conforming plain-text answer is the last assistant message at + // this point; record an explicit termination message so + // GetLastAssistantMessageContent() — what parents, MCP, and the API read + // as the final answer — reports the failure instead of that text. + addAgentMessage(sess, a, &chat.Message{ + Role: chat.MessageRoleAssistant, + Content: errMsg, + CreatedAt: r.now().Format(time.RFC3339), + }, events) + return turnExit, turnEndReasonError +} + +// findToolByName returns the tool with the given name from agentTools. +func findToolByName(agentTools []tools.Tool, name string) (tools.Tool, bool) { + for _, t := range agentTools { + if t.Name == name { + return t, true + } + } + return tools.Tool{}, false +} diff --git a/pkg/runtime/structured_output_test.go b/pkg/runtime/structured_output_test.go new file mode 100644 index 000000000..62d4c28e6 --- /dev/null +++ b/pkg/runtime/structured_output_test.go @@ -0,0 +1,725 @@ +package runtime + +import ( + "bytes" + "context" + "encoding/json" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/hooks" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/skills" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/tools" + skillstool "github.com/docker/docker-agent/pkg/tools/builtin/skills" + "github.com/docker/docker-agent/pkg/tools/builtin/structuredoutput" +) + +// toolModeStructuredOutput returns a tool-mode structured-output config with +// a single required string property "answer". +func toolModeStructuredOutput() *latest.StructuredOutput { + return &latest.StructuredOutput{ + Name: "answer_format", + Description: "The final answer", + Mode: latest.StructuredOutputModeTool, + Schema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "answer": map[string]any{"type": "string"}, + }, + "required": []any{"answer"}, + "additionalProperties": false, + }, + } +} + +// soRecordingProvider wraps queueProvider and captures the messages of every +// model call so tests can assert on transient system reminders. +type soRecordingProvider struct { + queueProvider + + mu sync.Mutex + allMsgs [][]chat.Message + allTools [][]tools.Tool +} + +func (p *soRecordingProvider) CreateChatCompletionStream(ctx context.Context, msgs []chat.Message, ts []tools.Tool) (chat.MessageStream, error) { + p.mu.Lock() + p.allMsgs = append(p.allMsgs, append([]chat.Message(nil), msgs...)) + p.allTools = append(p.allTools, append([]tools.Tool(nil), ts...)) + p.mu.Unlock() + return p.queueProvider.CreateChatCompletionStream(ctx, msgs, ts) +} + +func (p *soRecordingProvider) calls() ([][]chat.Message, [][]tools.Tool) { + p.mu.Lock() + defer p.mu.Unlock() + return p.allMsgs, p.allTools +} + +// structuredOutputRuntime builds a single-agent runtime whose model replays +// the given streams and whose agent uses the supplied structured output. +func structuredOutputRuntime(t *testing.T, so *latest.StructuredOutput, agentOpts []agent.Opt, streams ...chat.MessageStream) (*LocalRuntime, *soRecordingProvider) { + t.Helper() + + prov := &soRecordingProvider{queueProvider: queueProvider{id: "test/mock-model", streams: streams}} + opts := append([]agent.Opt{ + agent.WithModel(prov), + agent.WithStructuredOutput(so), + }, agentOpts...) + root := agent.New("root", "You are a test agent", opts...) + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + return rt, prov +} + +func runAndCollect(t *testing.T, rt *LocalRuntime, sess *session.Session) []Event { + t.Helper() + var events []Event + for ev := range rt.RunStream(t.Context(), sess) { + events = append(events, ev) + } + return events +} + +// outputCallStream scripts a model turn that calls the internal output tool +// with the given arguments and nothing else. +func outputCallStream(callID, args string) *mockStream { + return newStreamBuilder(). + AddToolCallName(callID, structuredoutput.ToolName). + AddToolCallArguments(callID, args). + AddToolCallStopWithUsage(5, 5). + Build() +} + +func toolResponsesByCallID(events []Event) map[string]*ToolCallResponseEvent { + responses := make(map[string]*ToolCallResponseEvent) + for _, ev := range events { + if resp, ok := ev.(*ToolCallResponseEvent); ok { + responses[resp.ToolCallID] = resp + } + } + return responses +} + +func errorEvents(events []Event) []*ErrorEvent { + var errs []*ErrorEvent + for _, ev := range events { + if errEv, ok := ev.(*ErrorEvent); ok { + errs = append(errs, errEv) + } + } + return errs +} + +// assertNoOrphanToolCalls verifies every tool call recorded in the session +// has a matching tool-role response, so the history stays provider-valid. +func assertNoOrphanToolCalls(t *testing.T, sess *session.Session) { + t.Helper() + responded := make(map[string]bool) + for _, m := range sess.GetAllMessages() { + if m.Message.Role == chat.MessageRoleTool { + responded[m.Message.ToolCallID] = true + } + } + for _, m := range sess.GetAllMessages() { + if m.Message.Role != chat.MessageRoleAssistant { + continue + } + for _, tc := range m.Message.ToolCalls { + assert.True(t, responded[tc.ID], "tool call %s has no tool response", tc.ID) + } + } +} + +// TestStructuredOutputToolMode_ExposesToolWithSchema proves the internal tool +// is offered to the model in tool mode — with Parameters exactly equal to the +// configured schema — and that native mode never exposes it. +func TestStructuredOutputToolMode_ExposesToolWithSchema(t *testing.T) { + t.Parallel() + + so := toolModeStructuredOutput() + rt, prov := structuredOutputRuntime(t, so, nil, + outputCallStream("call_1", `{"answer":"hi"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + runAndCollect(t, rt, sess) + + _, allTools := prov.calls() + require.NotEmpty(t, allTools) + var found *tools.Tool + for i := range allTools[0] { + if allTools[0][i].Name == structuredoutput.ToolName { + found = &allTools[0][i] + } + } + require.NotNil(t, found, "tool mode must offer the internal output tool to the model") + assert.Equal(t, so.Schema, found.Parameters, "tool parameters must be exactly the configured schema") +} + +// TestStructuredOutputNativeMode_DoesNotExposeTool pins the no-change +// contract for native (and absent) mode: no internal tool, plain text final. +func TestStructuredOutputNativeMode_DoesNotExposeTool(t *testing.T) { + t.Parallel() + + native := &latest.StructuredOutput{ + Name: "answer_format", + Schema: map[string]any{"type": "object"}, + } + rt, prov := structuredOutputRuntime(t, native, nil, + newStreamBuilder().AddContent("plain text").AddStopWithUsage(3, 2).Build(), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + assert.Equal(t, "plain text", sess.GetLastAssistantMessageContent()) + + _, allTools := prov.calls() + require.NotEmpty(t, allTools) + for _, tl := range allTools[0] { + assert.NotEqual(t, structuredoutput.ToolName, tl.Name, "native mode must not expose the internal tool") + } +} + +// TestStructuredOutputToolMode_ValidCallIsTerminal drives the happy path: a +// single valid output call records a tool response, appends a final +// assistant message whose content is the canonical JSON, and stops the run. +func TestStructuredOutputToolMode_ValidCallIsTerminal(t *testing.T) { + t.Parallel() + + rt, prov := structuredOutputRuntime(t, toolModeStructuredOutput(), nil, + outputCallStream("call_1", "{\n \"answer\": \"hi\" \n}"), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + last := sess.GetLastAssistantMessageContent() + assert.JSONEq(t, `{"answer":"hi"}`, last, + "final assistant message must carry the validated JSON") + // Canonicalization contract: the accepted JSON is stored compacted. + var compacted bytes.Buffer + require.NoError(t, json.Compact(&compacted, []byte(last))) + assert.Equal(t, compacted.String(), last, "final content must be canonical compact JSON") + + responses := toolResponsesByCallID(events) + require.Contains(t, responses, "call_1") + assert.False(t, responses["call_1"].Result.IsError) + + assertNoOrphanToolCalls(t, sess) + + // Terminal: exactly one model call, no retry iteration. + allMsgs, _ := prov.calls() + assert.Len(t, allMsgs, 1) + + // The synthetic final answer must not itself carry tool calls. + msgs := sess.GetAllMessages() + lastMsg := msgs[len(msgs)-1].Message + assert.Equal(t, chat.MessageRoleAssistant, lastMsg.Role) + assert.Empty(t, lastMsg.ToolCalls) +} + +// TestStructuredOutputToolMode_InvalidJSONAllowsCorrection verifies a +// schema-violating call yields an IsError tool response (with details) and +// the model gets another iteration to fix it. +func TestStructuredOutputToolMode_InvalidJSONAllowsCorrection(t *testing.T) { + t.Parallel() + + rt, prov := structuredOutputRuntime(t, toolModeStructuredOutput(), nil, + outputCallStream("call_1", `{"answer":42}`), + outputCallStream("call_2", `{"answer":"fixed"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + + responses := toolResponsesByCallID(events) + require.Contains(t, responses, "call_1") + require.NotNil(t, responses["call_1"].Result) + assert.True(t, responses["call_1"].Result.IsError, "schema violation must surface as a tool error") + assert.Contains(t, responses["call_1"].Response, "answer") + + require.Contains(t, responses, "call_2") + assert.False(t, responses["call_2"].Result.IsError) + + assert.JSONEq(t, `{"answer":"fixed"}`, sess.GetLastAssistantMessageContent()) + assertNoOrphanToolCalls(t, sess) + + allMsgs, _ := prov.calls() + assert.Len(t, allMsgs, 2, "the model must get a correction iteration") +} + +// TestStructuredOutputToolMode_MixedBatchNeverTerminal drives a batch mixing +// the output tool with a real tool: the output call is rejected as +// non-exclusive, the other tool runs normally, and only a later exclusive +// call finalizes the run. +func TestStructuredOutputToolMode_MixedBatchNeverTerminal(t *testing.T) { + t.Parallel() + + var noopRan bool + noop := tools.Tool{ + Name: "noop", + Parameters: map[string]any{}, + Handler: func(context.Context, tools.ToolCall, tools.Runtime) (*tools.ToolCallResult, error) { + noopRan = true + return tools.ResultSuccess("ok"), nil + }, + } + + mixed := newStreamBuilder(). + AddToolCallName("call_out", structuredoutput.ToolName). + AddToolCallArguments("call_out", `{"answer":"too early"}`). + AddToolCallName("call_noop", "noop"). + AddToolCallArguments("call_noop", `{}`). + AddToolCallStopWithUsage(5, 5). + Build() + + rt, _ := structuredOutputRuntime(t, toolModeStructuredOutput(), + []agent.Opt{agent.WithToolSets(newStubToolSet(nil, []tools.Tool{noop}, nil))}, + mixed, + outputCallStream("call_final", `{"answer":"done"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me"), session.WithToolsApproved(true)) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + + responses := toolResponsesByCallID(events) + require.Contains(t, responses, "call_out") + require.NotNil(t, responses["call_out"].Result) + assert.True(t, responses["call_out"].Result.IsError, "output call in a mixed batch must be rejected") + assert.Contains(t, responses["call_out"].Response, "only tool call") + + require.Contains(t, responses, "call_noop") + assert.False(t, responses["call_noop"].Result.IsError, "other tools in the batch keep their normal dispatch") + assert.True(t, noopRan) + + assert.JSONEq(t, `{"answer":"done"}`, sess.GetLastAssistantMessageContent(), + "only the later exclusive call may produce the terminal result") + assertNoOrphanToolCalls(t, sess) +} + +// TestStructuredOutputToolMode_MultiOutputBatchNeverTerminal rejects every +// output call when several arrive in one batch, even when each would +// validate on its own. +func TestStructuredOutputToolMode_MultiOutputBatchNeverTerminal(t *testing.T) { + t.Parallel() + + multi := newStreamBuilder(). + AddToolCallName("call_a", structuredoutput.ToolName). + AddToolCallArguments("call_a", `{"answer":"first"}`). + AddToolCallName("call_b", structuredoutput.ToolName). + AddToolCallArguments("call_b", `{"answer":"second"}`). + AddToolCallStopWithUsage(5, 5). + Build() + + rt, _ := structuredOutputRuntime(t, toolModeStructuredOutput(), nil, + multi, + outputCallStream("call_final", `{"answer":"done"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + + responses := toolResponsesByCallID(events) + for _, id := range []string{"call_a", "call_b"} { + require.Contains(t, responses, id) + require.NotNil(t, responses[id].Result) + assert.True(t, responses[id].Result.IsError, "call %s must be rejected as non-exclusive", id) + } + assert.JSONEq(t, `{"answer":"done"}`, sess.GetLastAssistantMessageContent()) + assertNoOrphanToolCalls(t, sess) +} + +// telemetryRuntime builds a single-agent tool-mode runtime with a +// recordingTelemetry injected, so tests can assert on tool-call records. +func telemetryRuntime(t *testing.T, streams ...chat.MessageStream) (*LocalRuntime, *recordingTelemetry) { + t.Helper() + + prov := &queueProvider{id: "test/mock-model", streams: streams} + root := agent.New("root", "You are a test agent", + agent.WithModel(prov), + agent.WithStructuredOutput(toolModeStructuredOutput()), + ) + rec := &recordingTelemetry{} + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), + WithSessionCompaction(false), WithModelStore(mockModelStore{}), WithTelemetry(rec)) + require.NoError(t, err) + return rt, rec +} + +// TestStructuredOutputToolMode_TelemetryPerOutputCall proves every +// intercepted output call records exactly one tool-call telemetry entry — +// the schema-invalid attempt as a failure, the valid retry as a success — +// with the same naming the dispatcher uses. +func TestStructuredOutputToolMode_TelemetryPerOutputCall(t *testing.T) { + t.Parallel() + + rt, rec := telemetryRuntime(t, + outputCallStream("call_1", `{"answer":42}`), + outputCallStream("call_2", `{"answer":"fixed"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + require.Empty(t, errorEvents(events)) + + calls := rec.snapshot().toolCalls + require.Len(t, calls, 2, "expected exactly one telemetry record per output call") + for _, c := range calls { + assert.Equal(t, structuredoutput.ToolName, c.ToolName) + assert.Equal(t, sess.ID, c.SessionID) + assert.Equal(t, "root", c.AgentName) + } + require.Error(t, calls[0].Err, "schema-invalid call must be recorded as a failure") + assert.Contains(t, calls[0].Err.Error(), "schema") + assert.NoError(t, calls[1].Err, "valid call must be recorded as a success") +} + +// TestStructuredOutputToolMode_TelemetryOnBatchRejection proves rejected +// non-exclusive output calls are still instrumented: each rejection records +// a failed tool-call entry, the later exclusive call a successful one. +func TestStructuredOutputToolMode_TelemetryOnBatchRejection(t *testing.T) { + t.Parallel() + + multi := newStreamBuilder(). + AddToolCallName("call_a", structuredoutput.ToolName). + AddToolCallArguments("call_a", `{"answer":"first"}`). + AddToolCallName("call_b", structuredoutput.ToolName). + AddToolCallArguments("call_b", `{"answer":"second"}`). + AddToolCallStopWithUsage(5, 5). + Build() + + rt, rec := telemetryRuntime(t, + multi, + outputCallStream("call_final", `{"answer":"done"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + require.Empty(t, errorEvents(events)) + + calls := rec.snapshot().toolCalls + require.Len(t, calls, 3, "two rejected calls plus the final exclusive call") + for _, c := range calls { + assert.Equal(t, structuredoutput.ToolName, c.ToolName) + assert.Equal(t, sess.ID, c.SessionID) + assert.Equal(t, "root", c.AgentName) + } + require.Error(t, calls[0].Err, "rejected call_a must be recorded as a failure") + require.Error(t, calls[1].Err, "rejected call_b must be recorded as a failure") + assert.Contains(t, calls[0].Err.Error(), "only tool call") + assert.NoError(t, calls[2].Err, "the exclusive retry must be recorded as a success") +} + +// TestStructuredOutputToolMode_TextStopTriggersRemindersThenError verifies a +// plain-text finish never becomes the final answer: the runtime injects a +// transient system reminder for up to two retries, then fails with the +// structured_output_failed error code. +func TestStructuredOutputToolMode_TextStopTriggersRemindersThenError(t *testing.T) { + t.Parallel() + + rt, prov := structuredOutputRuntime(t, toolModeStructuredOutput(), nil, + newStreamBuilder().AddContent("text one").AddStopWithUsage(3, 2).Build(), + newStreamBuilder().AddContent("text two").AddStopWithUsage(3, 2).Build(), + newStreamBuilder().AddContent("text three").AddStopWithUsage(3, 2).Build(), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + errs := errorEvents(events) + require.Len(t, errs, 1, "exhausted reminders must produce exactly one coded error") + assert.Equal(t, ErrorCodeStructuredOutputFailed, errs[0].Code) + assert.Contains(t, errs[0].Error, structuredoutput.ToolName) + + // The explicit termination message is the final assistant message, so + // consumers of GetLastAssistantMessageContent() read the failure, never + // the non-conforming plain text. + last := sess.GetLastAssistantMessageContent() + assert.Contains(t, last, "did not deliver structured output") + assert.Contains(t, last, structuredoutput.ToolName) + assert.NotContains(t, last, "text three") + + allMsgs, _ := prov.calls() + require.Len(t, allMsgs, 3, "two reminders allow exactly three model calls") + + // The reminder is a transient system message: present in the retry + // calls, never persisted to the session. + for i, call := range allMsgs[1:] { + var reminded bool + for _, m := range call { + if m.Role == chat.MessageRoleSystem && strings.Contains(m.Content, structuredoutput.ToolName) { + reminded = true + } + } + assert.True(t, reminded, "retry call %d must carry the transient reminder", i+1) + } + var firstCallReminded bool + for _, m := range allMsgs[0] { + if m.Role == chat.MessageRoleSystem && strings.Contains(m.Content, "Reminder: this conversation requires structured output") { + firstCallReminded = true + } + } + assert.False(t, firstCallReminded, "the first call must not carry a reminder") + for _, m := range sess.GetAllMessages() { + assert.NotContains(t, m.Message.Content, "Reminder: this conversation requires structured output", + "the reminder must never be persisted to the session") + } +} + +// TestStructuredOutputToolMode_SuccessAfterReminder proves recovery: a text +// stop followed by a valid exclusive output call ends the run cleanly. +func TestStructuredOutputToolMode_SuccessAfterReminder(t *testing.T) { + t.Parallel() + + rt, prov := structuredOutputRuntime(t, toolModeStructuredOutput(), nil, + newStreamBuilder().AddContent("oops, plain text").AddStopWithUsage(3, 2).Build(), + outputCallStream("call_1", `{"answer":"recovered"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + assert.JSONEq(t, `{"answer":"recovered"}`, sess.GetLastAssistantMessageContent()) + + allMsgs, _ := prov.calls() + assert.Len(t, allMsgs, 2) + assertNoOrphanToolCalls(t, sess) +} + +// TestStructuredOutputToolMode_OtherToolTurnsDoNotConsumeReminders verifies +// ordinary tool-using turns keep working in tool mode and don't burn the +// reminder budget: tool turn → text stop (reminder 1) → tool turn → text +// stop (reminder 2) → valid output call still succeeds. +func TestStructuredOutputToolMode_OtherToolTurnsDoNotConsumeReminders(t *testing.T) { + t.Parallel() + + noop := tools.Tool{ + Name: "noop", + Parameters: map[string]any{}, + Handler: func(context.Context, tools.ToolCall, tools.Runtime) (*tools.ToolCallResult, error) { + return tools.ResultSuccess("ok"), nil + }, + } + noopStream := func(id string) chat.MessageStream { + return newStreamBuilder(). + AddToolCallName(id, "noop"). + AddToolCallArguments(id, `{}`). + AddToolCallStopWithUsage(2, 2). + Build() + } + + rt, prov := structuredOutputRuntime(t, toolModeStructuredOutput(), + []agent.Opt{agent.WithToolSets(newStubToolSet(nil, []tools.Tool{noop}, nil))}, + noopStream("call_n1"), + newStreamBuilder().AddContent("text one").AddStopWithUsage(3, 2).Build(), + noopStream("call_n2"), + newStreamBuilder().AddContent("text two").AddStopWithUsage(3, 2).Build(), + outputCallStream("call_final", `{"answer":"done"}`), + ) + + sess := session.New(session.WithUserMessage("Answer me"), session.WithToolsApproved(true)) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events), "tool turns must not consume the reminder budget") + assert.JSONEq(t, `{"answer":"done"}`, sess.GetLastAssistantMessageContent()) + + allMsgs, _ := prov.calls() + assert.Len(t, allMsgs, 5) + assertNoOrphanToolCalls(t, sess) +} + +// TestStructuredOutputToolMode_StopHooksReceiveJSON pins the stop-hook +// contract: finalization fires stop hooks with the validated JSON, exactly +// like a natural stop does with its text content. +func TestStructuredOutputToolMode_StopHooksReceiveJSON(t *testing.T) { + t.Parallel() + + prov := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{ + outputCallStream("call_1", `{"answer":"hooked"}`), + }} + root := agent.New("root", "You are a test agent", + agent.WithModel(prov), + agent.WithStructuredOutput(toolModeStructuredOutput()), + agent.WithHooks(&hooks.Config{ + Stop: []hooks.Hook{{Type: hooks.HookTypeBuiltin, Command: "test_record_stop"}}, + }), + ) + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + + rb := &recordingBuiltin{} + require.NoError(t, rt.hooksRegistry.RegisterBuiltin("test_record_stop", rb.hook)) + rt.buildHooksExecutors() + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + inputs := rb.snapshot() + require.NotEmpty(t, inputs, "stop hooks must fire on structured-output finalization") + assert.JSONEq(t, `{"answer":"hooked"}`, inputs[len(inputs)-1].StopResponse) +} + +// TestStructuredOutputToolMode_ForceHandoffAfterFinalization preserves the +// force_handoff semantics: after the structured result finalizes the turn, +// the conversation is routed to the configured target agent. +func TestStructuredOutputToolMode_ForceHandoffAfterFinalization(t *testing.T) { + t.Parallel() + + sumProv := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{ + newStreamBuilder().AddContent("summary of the JSON").AddStopWithUsage(2, 2).Build(), + }} + summarizer := agent.New("summarizer", "You summarize", agent.WithModel(sumProv)) + + rootProv := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{ + outputCallStream("call_1", `{"answer":"handed off"}`), + }} + root := agent.New("root", "You extract", + agent.WithModel(rootProv), + agent.WithStructuredOutput(toolModeStructuredOutput()), + agent.WithForceHandoff(summarizer), + ) + + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root, summarizer)), WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + + sess := session.New(session.WithUserMessage("Answer me")) + sess.Title = "Unit Test" + events := runAndCollect(t, rt, sess) + + require.Empty(t, errorEvents(events)) + assert.Equal(t, "summarizer", rt.CurrentAgentName(t.Context())) + assert.Equal(t, "summary of the JSON", sess.GetLastAssistantMessageContent()) + + // The structured result is still in the transcript, before the handoff. + var sawJSON bool + for _, m := range sess.GetAllMessages() { + if m.Message.Role == chat.MessageRoleAssistant && m.Message.Content == `{"answer":"handed off"}` { + sawJSON = true + } + } + assert.True(t, sawJSON) +} + +// TestStructuredOutputToolMode_ReservedNameCollisionFails pins the collision +// contract: a real tool that claims the reserved name fails the run loudly +// instead of masking either tool. +func TestStructuredOutputToolMode_ReservedNameCollisionFails(t *testing.T) { + t.Parallel() + + impostor := tools.Tool{ + Name: structuredoutput.ToolName, + Parameters: map[string]any{}, + Handler: func(context.Context, tools.ToolCall, tools.Runtime) (*tools.ToolCallResult, error) { + return tools.ResultSuccess("impostor"), nil + }, + } + + rt, _ := structuredOutputRuntime(t, toolModeStructuredOutput(), + []agent.Opt{agent.WithToolSets(newStubToolSet(nil, []tools.Tool{impostor}, nil))}, + ) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + errs := errorEvents(events) + require.NotEmpty(t, errs) + assert.Equal(t, ErrorCodeToolFailed, errs[0].Code) + assert.Contains(t, errs[0].Error, structuredoutput.ToolName) + assert.Contains(t, errs[0].Error, "reserved") +} + +// TestStructuredOutputToolMode_ForkSkillChildIsExempt proves fork-mode skill +// children of a tool-mode agent are exempt from enforcement: with a skill +// allowed-tools list that excludes the internal tool, a plain-text answer +// succeeds without reminders or errors, and the child session carries +// DisableStructuredOutput. +func TestStructuredOutputToolMode_ForkSkillChildIsExempt(t *testing.T) { + t.Parallel() + + skillTS := skillstool.New([]skills.Skill{{ + Name: "greet", + Description: "Greets the user", + Context: "fork", + InlineContent: "# Greet\nSay the greeting.", + AllowedTools: []string{"noop"}, // excludes __structured_output__ + }}, "") + + rt, prov := structuredOutputRuntime(t, toolModeStructuredOutput(), + []agent.Opt{agent.WithToolSets(skillTS)}, + newStreamBuilder().AddContent("greeting delivered").AddStopWithUsage(3, 2).Build(), + ) + + sess := session.New(session.WithUserMessage("Run the greet skill")) + evts := make(chan Event, 256) + result, err := rt.RunSkillFork(t.Context(), sess, + skillstool.RunSkillArgs{Name: "greet", Task: "greet the user"}, NewChannelSink(evts)) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.IsError, "the skill child's plain-text answer must be accepted: %s", result.Output) + assert.Equal(t, "greeting delivered", result.Output) + + close(evts) + for ev := range evts { + if errEv, ok := ev.(*ErrorEvent); ok { + t.Errorf("unexpected error event: %s", errEv.Error) + } + } + + child := firstSubSession(sess) + require.NotNil(t, child) + assert.True(t, child.DisableStructuredOutput, + "the skill child must be exempt from structured-output enforcement") + + allMsgs, allTools := prov.calls() + require.Len(t, allMsgs, 1, "the plain-text stop is final: no reminder retry") + for _, tl := range allTools[0] { + assert.NotEqual(t, structuredoutput.ToolName, tl.Name, + "the internal output tool must not be offered to the skill child") + } +} + +// TestStructuredOutputToolMode_ProgrammaticBrokenSchemaFailsRun proves the +// compile error of a programmatically-built agent (no teamloader fail-fast) +// still surfaces: the run fails at tool collection, before any model call. +func TestStructuredOutputToolMode_ProgrammaticBrokenSchemaFailsRun(t *testing.T) { + t.Parallel() + + broken := &latest.StructuredOutput{ + Name: "broken", + Mode: latest.StructuredOutputModeTool, + Schema: map[string]any{"type": 42}, + } + rt, prov := structuredOutputRuntime(t, broken, nil) + + sess := session.New(session.WithUserMessage("Answer me")) + events := runAndCollect(t, rt, sess) + + errs := errorEvents(events) + require.NotEmpty(t, errs) + assert.Equal(t, ErrorCodeToolFailed, errs[0].Code) + assert.Contains(t, errs[0].Error, "structured_output") + + allMsgs, _ := prov.calls() + assert.Empty(t, allMsgs, "a broken schema must fail the run before any model call") +} diff --git a/pkg/session/branch.go b/pkg/session/branch.go index ac639eaea..672aa3e73 100644 --- a/pkg/session/branch.go +++ b/pkg/session/branch.go @@ -107,6 +107,7 @@ func (s *Session) Clone() *Session { ExcludedTools: cloneStringSlice(s.ExcludedTools), AllowedTools: cloneStringSlice(s.AllowedTools), ExtraToolSets: slices.Clone(s.ExtraToolSets), + DisableStructuredOutput: s.DisableStructuredOutput, AgentName: s.AgentName, ParentID: s.ParentID, DelegationLineage: cloneStringSlice(s.DelegationLineage), @@ -206,10 +207,11 @@ func copySessionMetadata(dst, src *Session, title string) { if src == nil || dst == nil { return } - // AgentName, ParentID, and DelegationLineage are deliberately not - // copied: a branch/fork is a fresh top-level session, not a pinned - // sub-session, so transient delegation state must not carry over. - // Clone() keeps them because it reproduces the session verbatim. + // AgentName, ParentID, DelegationLineage, and DisableStructuredOutput + // are deliberately not copied: a branch/fork is a fresh top-level + // session, not a pinned sub-session, so transient delegation state must + // not carry over. Clone() keeps them because it reproduces the session + // verbatim. dst.SetTitle(title) dst.ToolsApproved = src.ToolsApproved dst.SafetyPolicy = src.SafetyPolicy diff --git a/pkg/session/clone_test.go b/pkg/session/clone_test.go index 36971e828..082bd7739 100644 --- a/pkg/session/clone_test.go +++ b/pkg/session/clone_test.go @@ -39,6 +39,7 @@ func TestClone_CopiesScalarFields(t *testing.T) { CustomModelsUsed: []string{"openai/gpt-4o"}, AttachedFiles: []string{"/abs/path.txt"}, ExcludedTools: []string{"run_skill"}, + DisableStructuredOutput: true, AgentName: "root", ParentID: "parent", DelegationLineage: []string{"coordinator"}, @@ -63,6 +64,7 @@ func TestClone_CopiesScalarFields(t *testing.T) { assert.Equal(t, int64(11), clone.InputTokens) assert.Equal(t, int64(22), clone.OutputTokens) assert.InEpsilon(t, 1.5, clone.Cost, 1e-9) + assert.True(t, clone.DisableStructuredOutput) assert.Equal(t, "root", clone.AgentName) assert.Equal(t, "parent", clone.ParentID) assert.Equal(t, []string{"coordinator"}, clone.DelegationLineage) diff --git a/pkg/session/session.go b/pkg/session/session.go index 09804298e..3382155bd 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -362,6 +362,15 @@ type Session struct { // filter (the skill explicitly asked for them). ExtraToolSets []tools.ToolSet `json:"-"` + // DisableStructuredOutput, when true, exempts this session from tool-mode + // structured-output enforcement: the internal output tool is not offered + // and plain-text finals are accepted. Used by fork-mode skill + // sub-sessions, whose answer is free-form text for the calling agent, not + // the agent's schema-constrained final output. Not persisted: like the + // other transient sub-session state, a restored or branched session + // starts with enforcement re-enabled. + DisableStructuredOutput bool `json:"-"` + // AgentName, when set, tells RunStream which agent to use for this session // instead of reading from the shared runtime currentAgent field. This is // required for background agent tasks where multiple sessions may run @@ -1541,6 +1550,15 @@ func WithExtraToolSets(toolSets []tools.ToolSet) Opt { } } +// WithStructuredOutputDisabled exempts the session from tool-mode +// structured-output enforcement. Used by fork-mode skill sub-sessions whose +// answers are free-form text for the calling agent. +func WithStructuredOutputDisabled(disabled bool) Opt { + return func(s *Session) { + s.DisableStructuredOutput = disabled + } +} + // WithAttachedFiles seeds the session with absolute paths of files the user // attached. Used when creating sub-sessions so that delegated agents inherit // the parent's file context. Empty and duplicate paths are dropped. diff --git a/pkg/session/session_options_test.go b/pkg/session/session_options_test.go index 5d0c4108f..935725ac2 100644 --- a/pkg/session/session_options_test.go +++ b/pkg/session/session_options_test.go @@ -24,6 +24,13 @@ func TestWithWorkingDir_EmptyReturnsNilAllowedDirs(t *testing.T) { assert.Nil(t, s.AllowedDirectories()) } +func TestWithStructuredOutputDisabled(t *testing.T) { + t.Parallel() + assert.False(t, New().DisableStructuredOutput) + assert.False(t, New(WithStructuredOutputDisabled(false)).DisableStructuredOutput) + assert.True(t, New(WithStructuredOutputDisabled(true)).DisableStructuredOutput) +} + func TestNewSession_AllOptionsApplied(t *testing.T) { t.Parallel() s := New( diff --git a/pkg/teamloader/teamloader.go b/pkg/teamloader/teamloader.go index 5aadbd1aa..3b780ce55 100644 --- a/pkg/teamloader/teamloader.go +++ b/pkg/teamloader/teamloader.go @@ -320,6 +320,7 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c agent.WithSessionCompaction(agentConfig.SessionCompactionEnabled()), agent.WithCommands(expander.ExpandCommands(ctx, agentConfig.Commands)), agent.WithHooks(config.MergeHooks(config.MergeHooks(agentConfig.Hooks, globalHooks), cliHooks)), + agent.WithStructuredOutput(agentConfig.StructuredOutput), } if agentConfig.Cache != nil && agentConfig.Cache.Enabled { @@ -423,6 +424,15 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c opts = append(opts, agent.WithToolSets(agentTools...)) ag := agent.New(agentConfig.Name, expander.Expand(ctx, agentConfig.Instruction, nil), opts...) + + // Tool-mode structured output needs a compilable JSON schema; catch a + // broken one at load time instead of on the first model turn. Going + // through the agent's lazy cache also pre-compiles the tool the + // runtime will reuse on every turn. + if _, err := ag.StructuredOutputTool(); err != nil { + return nil, fmt.Errorf("agent %s: structured_output: %w", agentConfig.Name, err) + } + agents = append(agents, ag) agentsByName[agentConfig.Name] = ag } diff --git a/pkg/teamloader/teamloader_test.go b/pkg/teamloader/teamloader_test.go index 6acb21618..54a50d40b 100644 --- a/pkg/teamloader/teamloader_test.go +++ b/pkg/teamloader/teamloader_test.go @@ -1133,6 +1133,62 @@ func TestLoadPropagatesMaxToolResultTokens(t *testing.T) { assert.Equal(t, 512, agt.MaxToolResultTokens()) } +// TestLoadPropagatesStructuredOutput verifies the original structured-output +// config (including tool mode) travels from the YAML to the built agent, +// where the runtime reads it via agent.StructuredOutput(). +func TestLoadPropagatesStructuredOutput(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "dummy") + + data := []byte(`agents: + root: + model: openai/gpt-4o + instruction: test + structured_output: + mode: tool + name: result + schema: + type: object + properties: + answer: + type: string + required: ["answer"] +`) + + team, err := Load(t.Context(), config.NewBytesSource("so.yaml", data), &config.RuntimeConfig{}, withTestProviderRegistry()...) + require.NoError(t, err) + + agt, err := team.Agent("root") + require.NoError(t, err) + so := agt.StructuredOutput() + require.NotNil(t, so, "agent must retain the original structured-output config") + assert.Equal(t, "result", so.Name) + assert.True(t, so.ToolMode()) + assert.Contains(t, so.Schema, "properties") +} + +// TestLoadRejectsUncompilableToolModeSchema pins the load-time failure for +// tool-mode structured output with a schema gojsonschema cannot compile. +func TestLoadRejectsUncompilableToolModeSchema(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "dummy") + + data := []byte(`agents: + root: + model: openai/gpt-4o + instruction: test + structured_output: + mode: tool + name: result + schema: + type: object + properties: + answer: + type: 42 +`) + + _, err := Load(t.Context(), config.NewBytesSource("bad-schema.yaml", data), &config.RuntimeConfig{}, withTestProviderRegistry()...) + require.ErrorContains(t, err, "agent root: structured_output") +} + // TestLoadPropagatesSafetyDefaults verifies the author-declared safety // defaults travel from the YAML config to the built team: runtime.safety // lands on the team (team.RuntimeSafety) and agents..safety on the diff --git a/pkg/tools/builtin/structuredoutput/structuredoutput.go b/pkg/tools/builtin/structuredoutput/structuredoutput.go new file mode 100644 index 000000000..866a5d03f --- /dev/null +++ b/pkg/tools/builtin/structuredoutput/structuredoutput.go @@ -0,0 +1,159 @@ +// Package structuredoutput implements the internal tool behind tool-mode +// structured output (structured_output.mode: tool). Instead of asking the +// provider for native structured output, the runtime exposes a single tool +// whose Parameters are exactly the configured JSON schema; the model +// delivers its final answer by calling it, and the runtime validates the +// arguments against the same schema before accepting the result. +package structuredoutput + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "strings" + + "github.com/xeipuuv/gojsonschema" + + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/tools" +) + +// ToolName is the reserved name of the internal structured-output tool. +// The dunder form is provider-compatible ([a-zA-Z0-9_-]) and unlikely to +// collide with real tools; the runtime fails loudly when a real tool +// claims it rather than masking either one. +const ToolName = "__structured_output__" + +// OutputTool bundles the structured-output configuration with its compiled +// JSON schema so validation and the tool definition stay in sync. +type OutputTool struct { + cfg *latest.StructuredOutput + schema *gojsonschema.Schema +} + +// New compiles cfg.Schema and returns the ready-to-expose tool. It fails +// when the configured schema is not a valid JSON Schema or references +// anything outside the document, so callers can surface the problem at +// load time instead of on the first model turn. +func New(cfg *latest.StructuredOutput) (*OutputTool, error) { + if cfg == nil { + return nil, errors.New("structured output configuration is required") + } + raw, err := json.Marshal(cfg.Schema) + if err != nil { + return nil, fmt.Errorf("marshaling schema: %w", err) + } + // Screen references on the marshaled form — exactly what gojsonschema + // will see — rather than cfg.Schema, whose nested values may use + // arbitrary Go types. + var decoded any + if err := json.Unmarshal(raw, &decoded); err != nil { + return nil, fmt.Errorf("decoding schema: %w", err) + } + if err := rejectNonLocalRefs(decoded); err != nil { + return nil, fmt.Errorf("invalid schema: %w", err) + } + schema, err := gojsonschema.NewSchema(gojsonschema.NewBytesLoader(raw)) + if err != nil { + return nil, fmt.Errorf("compiling schema: %w", err) + } + return &OutputTool{cfg: cfg, schema: schema}, nil +} + +// rejectNonLocalRefs walks a decoded JSON schema and rejects any $ref that +// is not a string starting with "#". gojsonschema resolves references at +// compile time, so an http(s)://, file:// or cross-document ref in an +// untrusted schema would trigger network or filesystem access (SSRF, +// local-file disclosure). Only same-document fragments are allowed. The +// walk is deliberately conservative: even a property that merely happens +// to be named "$ref" is rejected rather than risk missing a real one. +func rejectNonLocalRefs(node any) error { + switch v := node.(type) { + case map[string]any: + for key, val := range v { + if key == "$ref" { + ref, ok := val.(string) + if !ok { + return fmt.Errorf("$ref must be a string, got %v", val) + } + if !strings.HasPrefix(ref, "#") { + return fmt.Errorf("non-local $ref %q: only same-document references starting with %q are allowed", ref, "#") + } + continue + } + if err := rejectNonLocalRefs(val); err != nil { + return err + } + } + case []any: + for _, item := range v { + if err := rejectNonLocalRefs(item); err != nil { + return err + } + } + } + return nil +} + +// Definition returns the tool definition offered to the model. Parameters +// are exactly the configured JSON schema. +func (t *OutputTool) Definition() tools.Tool { + description := fmt.Sprintf( + "Deliver the final answer of this conversation as structured output (%s). "+ + "The arguments must be a JSON object matching the tool's parameter schema. "+ + "Call this tool alone, with no other tool calls in the same response; a valid call ends the turn.", + t.cfg.Name) + if t.cfg.Description != "" { + description += " Expected content: " + t.cfg.Description + } + return tools.Tool{ + Name: ToolName, + Category: "structured_output", + Description: description, + Parameters: t.cfg.Schema, + Handler: t.handle, + Annotations: tools.ToolAnnotations{ + ReadOnlyHint: true, + Title: "Structured Output", + }, + } +} + +// Validate checks rawJSON against the configured schema and returns its +// compacted, canonical form. The returned error is model-facing: it details +// what failed so the model can correct itself on the next attempt. +func (t *OutputTool) Validate(rawJSON string) (string, error) { + if strings.TrimSpace(rawJSON) == "" { + rawJSON = "{}" + } + var compacted bytes.Buffer + if err := json.Compact(&compacted, []byte(rawJSON)); err != nil { + return "", fmt.Errorf("arguments are not valid JSON: %w", err) + } + result, err := t.schema.Validate(gojsonschema.NewStringLoader(compacted.String())) + if err != nil { + return "", fmt.Errorf("validating arguments: %w", err) + } + if !result.Valid() { + details := make([]string, 0, len(result.Errors())) + for _, e := range result.Errors() { + details = append(details, e.String()) + } + return "", fmt.Errorf("arguments do not match the %q schema: %s", + t.cfg.Name, strings.Join(details, "; ")) + } + return compacted.String(), nil +} + +// handle is the pure tool handler: it validates the call arguments and +// returns either the canonical JSON as a success result or a detailed +// error result the model can act on. Terminality is the runtime's call. +func (t *OutputTool) handle(_ context.Context, toolCall tools.ToolCall, _ tools.Runtime) (*tools.ToolCallResult, error) { + out, err := t.Validate(toolCall.Function.Arguments) + if err != nil { + return tools.ResultError(fmt.Sprintf("Structured output rejected: %v. Fix the arguments and call %s again.", err, ToolName)), nil + } + return tools.ResultSuccess(out), nil +} diff --git a/pkg/tools/builtin/structuredoutput/structuredoutput_test.go b/pkg/tools/builtin/structuredoutput/structuredoutput_test.go new file mode 100644 index 000000000..cd3061ddc --- /dev/null +++ b/pkg/tools/builtin/structuredoutput/structuredoutput_test.go @@ -0,0 +1,231 @@ +package structuredoutput + +import ( + "bytes" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/tools" +) + +func personConfig() *latest.StructuredOutput { + return &latest.StructuredOutput{ + Name: "person_info", + Description: "Information about a person", + Mode: latest.StructuredOutputModeTool, + Schema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string"}, + "age": map[string]any{"type": "integer"}, + }, + "required": []any{"name", "age"}, + "additionalProperties": false, + }, + } +} + +func TestNew_RejectsNilAndBadSchema(t *testing.T) { + t.Parallel() + + _, err := New(nil) + require.Error(t, err) + + _, err = New(&latest.StructuredOutput{ + Name: "broken", + Schema: map[string]any{"type": 42}, + }) + require.ErrorContains(t, err, "schema") +} + +// TestNew_OnlyLocalRefsAllowed pins the security contract: gojsonschema +// resolves $ref at compile time, so anything but a same-document fragment +// (network, filesystem or cross-document access) must be rejected before +// compilation ever sees the schema. +func TestNew_OnlyLocalRefsAllowed(t *testing.T) { + t.Parallel() + + schemaWithRef := func(ref any) map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "a": map[string]any{"$ref": ref}, + }, + } + } + + tests := []struct { + name string + ref any + wantErr string + }{ + { + name: "http ref rejected", + ref: "http://127.0.0.1:1/schema.json", + wantErr: "non-local $ref", + }, + { + name: "file ref rejected", + ref: "file:///etc/passwd", + wantErr: "non-local $ref", + }, + { + name: "relative document ref rejected", + ref: "other.json#", + wantErr: "non-local $ref", + }, + { + name: "non-string ref rejected", + ref: 42, + wantErr: "$ref must be a string", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + _, err := New(&latest.StructuredOutput{Name: "refs", Schema: schemaWithRef(tt.ref)}) + require.ErrorContains(t, err, tt.wantErr) + }) + } +} + +// TestNew_LocalRefIsAcceptedAndResolved proves same-document fragments stay +// usable: a #/$defs/... ref compiles and enforces the referenced subschema, +// including when nested inside arrays (allOf). +func TestNew_LocalRefIsAcceptedAndResolved(t *testing.T) { + t.Parallel() + + ot, err := New(&latest.StructuredOutput{ + Name: "local_refs", + Schema: map[string]any{ + "type": "object", + "$defs": map[string]any{"answer": map[string]any{"type": "string"}}, + "properties": map[string]any{ + "answer": map[string]any{ + "allOf": []any{map[string]any{"$ref": "#/$defs/answer"}}, + }, + }, + "required": []any{"answer"}, + }, + }) + require.NoError(t, err) + + got, err := ot.Validate(`{"answer":"hi"}`) + require.NoError(t, err) + assert.JSONEq(t, `{"answer":"hi"}`, got) + + _, err = ot.Validate(`{"answer":42}`) + require.ErrorContains(t, err, "local_refs") +} + +func TestDefinition_UsesConfiguredSchemaAsParameters(t *testing.T) { + t.Parallel() + + cfg := personConfig() + ot, err := New(cfg) + require.NoError(t, err) + + def := ot.Definition() + assert.Equal(t, ToolName, def.Name) + assert.Equal(t, "structured_output", def.Category) + assert.Contains(t, def.Description, cfg.Name) + assert.Contains(t, def.Description, cfg.Description) + require.NotNil(t, def.Handler) + + params, ok := def.Parameters.(map[string]any) + require.True(t, ok, "Parameters must be exactly the configured schema object") + assert.Equal(t, cfg.Schema, params) +} + +func TestValidate(t *testing.T) { + t.Parallel() + + ot, err := New(personConfig()) + require.NoError(t, err) + + tests := []struct { + name string + in string + want string + wantErr string + }{ + { + name: "valid compacted", + in: `{"name":"Ada","age":36}`, + want: `{"name":"Ada","age":36}`, + }, + { + name: "valid is canonicalized", + in: "{\n \"name\": \"Ada\",\n \"age\": 36\n}", + want: `{"name":"Ada","age":36}`, + }, + { + name: "empty arguments fail required", + in: "", + wantErr: "person_info", + }, + { + name: "invalid JSON", + in: `{"name": "Ada",`, + wantErr: "not valid JSON", + }, + { + name: "schema violation names the field", + in: `{"name":"Ada","age":"old"}`, + wantErr: "age", + }, + { + name: "additional property rejected", + in: `{"name":"Ada","age":36,"extra":true}`, + wantErr: "person_info", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got, err := ot.Validate(tt.in) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestHandler_ReturnsErrorResultInsteadOfError(t *testing.T) { + t.Parallel() + + ot, err := New(personConfig()) + require.NoError(t, err) + def := ot.Definition() + + res, err := def.Handler(t.Context(), tools.ToolCall{ + ID: "call_1", + Function: tools.FunctionCall{Name: ToolName, Arguments: `{"name":"Ada"}`}, + }, tools.NopRuntime{}) + require.NoError(t, err) + require.NotNil(t, res) + assert.True(t, res.IsError) + assert.Contains(t, res.Output, "Structured output rejected") + assert.Contains(t, res.Output, ToolName) + + res, err = def.Handler(t.Context(), tools.ToolCall{ + ID: "call_2", + Function: tools.FunctionCall{Name: ToolName, Arguments: ` {"name":"Ada", "age": 36} `}, + }, tools.NopRuntime{}) + require.NoError(t, err) + require.NotNil(t, res) + assert.False(t, res.IsError) + assert.JSONEq(t, `{"name":"Ada","age":36}`, res.Output) + var compacted bytes.Buffer + require.NoError(t, json.Compact(&compacted, []byte(res.Output))) + assert.Equal(t, compacted.String(), res.Output, "handler output must be compacted") +}