From b023b0fd0bfa0e996463a7b9ba27e6ce674665af Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Thu, 9 Jul 2026 16:55:33 -0400 Subject: [PATCH 1/2] feat(sdk/go): Agent.Pause with webhook-resumed approvals Parity with the Python SDK's Agent.pause(): PauseManager registers a pending approval before client.RequestApproval transitions the execution to waiting, then blocks until the control plane's /webhooks/approval callback resolves it (or expiry/cancellation). Route matches the Python agent_server path and the CP's notifyApprovalCallback payload; exempted from origin/DID middleware like other CP-to-worker notifications. PauseClock intentionally not ported: the Go SDK has no execution watchdog to discount. Co-Authored-By: Claude Fable 5 --- sdk/go/agent/agent.go | 24 ++ sdk/go/agent/agent_lifecycle.go | 6 + sdk/go/agent/pause.go | 381 +++++++++++++++++++ sdk/go/agent/pause_test.go | 637 ++++++++++++++++++++++++++++++++ 4 files changed, 1048 insertions(+) create mode 100644 sdk/go/agent/pause.go create mode 100644 sdk/go/agent/pause_test.go diff --git a/sdk/go/agent/agent.go b/sdk/go/agent/agent.go index 4806de7a7..872b64aa8 100644 --- a/sdk/go/agent/agent.go +++ b/sdk/go/agent/agent.go @@ -517,6 +517,11 @@ type Agent struct { cancelMu sync.Mutex cancelFuncs map[string]context.CancelFunc + // pauseManager tracks pending Agent.Pause() calls, keyed by + // approval_request_id, and resolves them when the control plane POSTs an + // approval resolution to /webhooks/approval. + pauseManager *PauseManager + startTime time.Time } @@ -576,6 +581,7 @@ func New(cfg Config) (*Agent, error) { logger: cfg.Logger, realtimeValidationFunctions: make(map[string]struct{}), cancelFuncs: make(map[string]context.CancelFunc), + pauseManager: NewPauseManager(), startTime: time.Now(), } @@ -821,6 +827,7 @@ func (a *Agent) handler() http.Handler { mux.HandleFunc("/reasoners/", a.handleReasoner) mux.HandleFunc("/skills/", a.handleSkill) mux.HandleFunc("/_internal/executions/", a.handleInternalCancel) + mux.HandleFunc("/webhooks/approval", a.handleApprovalWebhook) var handler http.Handler = mux @@ -852,6 +859,15 @@ func (a *Agent) originAuthMiddleware(next http.Handler, token string) http.Handl next.ServeHTTP(w, r) return } + // /webhooks/approval is a control-plane→worker approval-resolution + // callback delivered unauthenticated (see notifyApprovalCallback in the + // control plane), analogous to the cancel notification. It only + // resolves a pause the agent itself initiated, so treat it as an open + // infrastructure route rather than a caller-initiated invocation. + if path == "/webhooks/approval" { + next.ServeHTTP(w, r) + return + } expected := "Bearer " + token if r.Header.Get("Authorization") != expected { @@ -883,6 +899,14 @@ func (a *Agent) localVerificationMiddleware(next http.Handler) http.Handler { next.ServeHTTP(w, r) return } + // /webhooks/approval is a control-plane→worker approval-resolution + // callback, delivered unauthenticated by the control plane. It carries + // no DID signature; skip local verification (same rationale as the + // cancel notification above). + if path == "/webhooks/approval" { + next.ServeHTTP(w, r) + return + } // Extract function name to check realtime validation requirement funcName := "" diff --git a/sdk/go/agent/agent_lifecycle.go b/sdk/go/agent/agent_lifecycle.go index 94a21831b..e9f4fae8a 100644 --- a/sdk/go/agent/agent_lifecycle.go +++ b/sdk/go/agent/agent_lifecycle.go @@ -340,6 +340,12 @@ func (a *Agent) shutdown(ctx context.Context) error { close(a.stopLease) } + // Unblock any reasoner still parked in Agent.Pause() so shutdown does not + // hang waiting on an approval callback that will never arrive. + if a.pauseManager != nil { + a.pauseManager.CancelAll() + } + if a.client != nil { if _, err := a.client.Shutdown(ctx, a.cfg.NodeID, types.ShutdownRequest{Reason: "shutdown", Version: a.cfg.Version}); err != nil { a.logger.Printf("failed to notify shutdown: %v", err) diff --git a/sdk/go/agent/pause.go b/sdk/go/agent/pause.go new file mode 100644 index 000000000..4b465c3a1 --- /dev/null +++ b/sdk/go/agent/pause.go @@ -0,0 +1,381 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "sync" + "time" + + "github.com/Agent-Field/agentfield/sdk/go/client" +) + +// Approval decision values carried on an ApprovalResult. They mirror the +// Python SDK's ApprovalResult.decision strings so a reasoner authored against +// either SDK reads the same values. +const ( + // ApprovalApproved is set when the human approved the request. + ApprovalApproved = "approved" + // ApprovalRejected is set when the human rejected the request. + ApprovalRejected = "rejected" + // ApprovalRequestChanges is set when the human asked for changes. + ApprovalRequestChanges = "request_changes" + // ApprovalExpired is set when the pause timed out without a resolution. + ApprovalExpired = "expired" + // ApprovalError is set when the control plane could not be notified. + ApprovalError = "error" + // ApprovalCancelled is set when the pause was cancelled (e.g. on shutdown). + ApprovalCancelled = "cancelled" +) + +// ApprovalResult is the outcome of a paused execution once a human (or an +// upstream service) resolves it. It mirrors the Python SDK's ApprovalResult +// dataclass field-for-field: decision, feedback, execution_id, +// approval_request_id and raw_response. +type ApprovalResult struct { + // Decision is one of the Approval* constants ("approved", "rejected", + // "request_changes", "expired", "error", "cancelled") or any custom value + // the control plane forwarded. + Decision string `json:"decision"` + // Feedback is a free-form human message accompanying the decision. + Feedback string `json:"feedback"` + // ExecutionID is the execution that was paused. + ExecutionID string `json:"execution_id"` + // ApprovalRequestID is the ID of the external approval request. + ApprovalRequestID string `json:"approval_request_id"` + // RawResponse is the parsed `response` payload from the callback, if any. + RawResponse map[string]any `json:"raw_response,omitempty"` +} + +// Approved reports whether the human approved the request. +func (r ApprovalResult) Approved() bool { return r.Decision == ApprovalApproved } + +// ChangesRequested reports whether the human asked for changes rather than +// approving or rejecting. +func (r ApprovalResult) ChangesRequested() bool { return r.Decision == ApprovalRequestChanges } + +// pendingPause holds the buffered channel that delivers the ApprovalResult to +// the blocked Pause caller. The channel is buffered (size 1) and closed after +// a single send so Resolve never blocks and a late Resolve cannot double-send. +type pendingPause struct { + ch chan ApprovalResult + resolved bool +} + +// PauseManager is a concurrency-safe registry of pending execution pauses, +// keyed by approval_request_id and resolved by the /webhooks/approval callback. +// It mirrors the Python SDK's _PauseManager and the TypeScript SDK's +// PauseManager. +type PauseManager struct { + mu sync.Mutex + pending map[string]*pendingPause + execToRequest map[string]string +} + +// NewPauseManager constructs an empty PauseManager. +func NewPauseManager() *PauseManager { + return &PauseManager{ + pending: make(map[string]*pendingPause), + execToRequest: make(map[string]string), + } +} + +// Register creates a pending pause for approvalRequestID and returns a +// receive-only channel that yields exactly one ApprovalResult when the pause +// resolves. Registration is idempotent: a second Register for the same +// approvalRequestID returns the existing channel rather than replacing it, so +// a fast callback that arrives before the caller blocks is never lost. +// executionID, when non-empty, is recorded for fallback resolution. +func (m *PauseManager) Register(approvalRequestID, executionID string) <-chan ApprovalResult { + m.mu.Lock() + defer m.mu.Unlock() + if existing, ok := m.pending[approvalRequestID]; ok { + return existing.ch + } + p := &pendingPause{ch: make(chan ApprovalResult, 1)} + m.pending[approvalRequestID] = p + if executionID != "" { + m.execToRequest[executionID] = approvalRequestID + } + return p.ch +} + +// Resolve resolves the pending pause identified by approvalRequestID with the +// given result. It returns true if a waiter was found and resolved, false +// otherwise (unknown or already-resolved request). +func (m *PauseManager) Resolve(approvalRequestID string, result ApprovalResult) bool { + m.mu.Lock() + defer m.mu.Unlock() + return m.resolveLocked(approvalRequestID, result) +} + +// resolveLocked resolves a pending pause; the caller must hold m.mu. +func (m *PauseManager) resolveLocked(approvalRequestID string, result ApprovalResult) bool { + p, ok := m.pending[approvalRequestID] + if !ok { + return false + } + delete(m.pending, approvalRequestID) + for eid, rid := range m.execToRequest { + if rid == approvalRequestID { + delete(m.execToRequest, eid) + break + } + } + if p.resolved { + return false + } + p.resolved = true + p.ch <- result // buffered (cap 1) — never blocks + close(p.ch) + return true +} + +// ResolveByExecutionID resolves a pending pause by execution_id, used as a +// fallback when the callback omits the approval_request_id. Returns true if a +// waiter was found. +func (m *PauseManager) ResolveByExecutionID(executionID string, result ApprovalResult) bool { + m.mu.Lock() + defer m.mu.Unlock() + rid, ok := m.execToRequest[executionID] + if !ok { + return false + } + return m.resolveLocked(rid, result) +} + +// CancelAll resolves every pending pause with a cancelled result. Used on +// shutdown so a reasoner blocked in Pause does not hang the process forever. +func (m *PauseManager) CancelAll() { + m.mu.Lock() + defer m.mu.Unlock() + for rid, p := range m.pending { + if !p.resolved { + p.resolved = true + p.ch <- ApprovalResult{ + Decision: ApprovalCancelled, + Feedback: "agent shutting down", + ApprovalRequestID: rid, + } + close(p.ch) + } + } + m.pending = make(map[string]*pendingPause) + m.execToRequest = make(map[string]string) +} + +// PendingCount returns the number of currently-pending pauses. Useful for tests. +func (m *PauseManager) PendingCount() int { + m.mu.Lock() + defer m.mu.Unlock() + return len(m.pending) +} + +// PauseOptions configures Agent.Pause. It mirrors the keyword arguments of the +// Python SDK's Agent.pause(). +type PauseOptions struct { + // ApprovalRequestID is the ID of the approval request the agent already + // created on an external service. Required. + ApprovalRequestID string + // ApprovalRequestURL is where a human can review the request. Optional. + ApprovalRequestURL string + // ExpiresInHours is passed to the control plane and, when Timeout is zero, + // also bounds how long Pause blocks. Defaults to 72 when <= 0. + ExpiresInHours int + // Timeout bounds how long Pause blocks. Zero defaults to + // ExpiresInHours hours. + Timeout time.Duration + // ExecutionID overrides the current execution. Defaults to the execution + // carried on ctx. + ExecutionID string +} + +// Pause pauses the current execution for external approval. +// +// It transitions the execution to "waiting" on the control plane (via +// client.RequestApproval), then blocks until the approval webhook callback +// resolves it, the context is cancelled, or the timeout elapses. The agent is +// responsible for creating the approval request on an external service and +// passing the resulting ApprovalRequestID. +// +// On timeout Pause returns an ApprovalResult with Decision=="expired" (not an +// error), mirroring the Python SDK. On context cancellation it returns the +// context error. If the control plane cannot be notified it returns an error +// after resolving the pending pause with Decision=="error". +// +// The agent must be serving (a reachable PublicURL) and connected to a control +// plane, since the callback URL is required for resolution. +func (a *Agent) Pause(ctx context.Context, opts PauseOptions) (*ApprovalResult, error) { + if ctx == nil { + ctx = context.Background() + } + + executionID := strings.TrimSpace(opts.ExecutionID) + if executionID == "" { + executionID = executionContextFrom(ctx).ExecutionID + } + if executionID == "" { + return nil, errors.New("no execution_id available — cannot pause") + } + if strings.TrimSpace(opts.ApprovalRequestID) == "" { + return nil, errors.New("approval_request_id is required — cannot pause") + } + if a.client == nil { + return nil, errors.New("agent is not connected to a control plane — cannot pause") + } + + // The callback URL must match the /webhooks/approval route registered on + // the agent's mux and be reachable from the control plane, so it is built + // from PublicURL (the same base URL advertised at registration). + base := strings.TrimSuffix(strings.TrimSpace(a.cfg.PublicURL), "/") + if base == "" { + return nil, errors.New("agent has no public URL — cannot build approval callback URL") + } + callbackURL := base + "/webhooks/approval" + + expiresInHours := opts.ExpiresInHours + if expiresInHours <= 0 { + expiresInHours = 72 + } + + // Register the pending pause BEFORE telling the control plane, so a fast + // callback that arrives before RequestApproval returns is not missed. + future := a.pauseManager.Register(opts.ApprovalRequestID, executionID) + + if _, err := a.client.RequestApproval(ctx, a.cfg.NodeID, executionID, client.RequestApprovalRequest{ + ApprovalRequestID: opts.ApprovalRequestID, + ApprovalRequestURL: opts.ApprovalRequestURL, + CallbackURL: callbackURL, + ExpiresInHours: expiresInHours, + }); err != nil { + // Clean up the pending pause if we could not even notify the CP. + a.pauseManager.Resolve(opts.ApprovalRequestID, ApprovalResult{ + Decision: ApprovalError, + Feedback: "failed to notify control plane", + ExecutionID: executionID, + ApprovalRequestID: opts.ApprovalRequestID, + }) + return nil, fmt.Errorf("pause: request approval: %w", err) + } + + a.Note(ctx, fmt.Sprintf("Execution paused — waiting for approval %s", opts.ApprovalRequestID), "approval", "waiting") + + timeout := opts.Timeout + if timeout <= 0 { + timeout = time.Duration(expiresInHours) * time.Hour + } + timer := time.NewTimer(timeout) + defer timer.Stop() + + select { + case result := <-future: + return &result, nil + case <-timer.C: + // Timeout is a normal outcome — return an "expired" result rather than + // erroring, matching the Python SDK. Resolve the pending pause so a + // late callback does not leak a resolved-but-unawaited entry. + expired := ApprovalResult{ + Decision: ApprovalExpired, + Feedback: "timed out waiting for approval", + ExecutionID: executionID, + ApprovalRequestID: opts.ApprovalRequestID, + } + a.pauseManager.Resolve(opts.ApprovalRequestID, expired) + return &expired, nil + case <-ctx.Done(): + // Cooperative cancellation — drop the pending pause and surface the + // context error to the caller (idiomatic Go). + a.pauseManager.Resolve(opts.ApprovalRequestID, ApprovalResult{ + Decision: ApprovalCancelled, + Feedback: "context cancelled", + ExecutionID: executionID, + ApprovalRequestID: opts.ApprovalRequestID, + }) + return nil, ctx.Err() + } +} + +// handleApprovalWebhook is the worker side of the approval callback transport. +// The control plane POSTs here (via the callback_url registered when the +// execution paused) once a human resolves the approval. The body carries +// {execution_id, decision, feedback, approval_request_id, response, new_status}; +// we resolve the matching pending pause — first by approval_request_id, then by +// execution_id as a fallback — and reply {status, resolved}. This matches the +// Python SDK's /webhooks/approval route and the TypeScript SDK's +// installApprovalWebhookRoute. +func (a *Agent) handleApprovalWebhook(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + defer r.Body.Close() + + var body struct { + ExecutionID string `json:"execution_id"` + Decision string `json:"decision"` + Feedback string `json:"feedback"` + ApprovalRequestID string `json:"approval_request_id"` + Response json.RawMessage `json:"response"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil && !errors.Is(err, io.EOF) { + writeJSON(w, http.StatusBadRequest, map[string]any{"error": "invalid JSON"}) + return + } + + if body.ExecutionID == "" || body.Decision == "" { + writeJSON(w, http.StatusBadRequest, map[string]any{ + "error": "execution_id and decision are required", + }) + return + } + + result := ApprovalResult{ + Decision: body.Decision, + Feedback: body.Feedback, + ExecutionID: body.ExecutionID, + ApprovalRequestID: body.ApprovalRequestID, + RawResponse: parseRawApprovalResponse(body.Response), + } + + resolved := false + if body.ApprovalRequestID != "" { + resolved = a.pauseManager.Resolve(body.ApprovalRequestID, result) + } + if !resolved && body.ExecutionID != "" { + resolved = a.pauseManager.ResolveByExecutionID(body.ExecutionID, result) + } + + writeJSON(w, http.StatusOK, map[string]any{"status": "received", "resolved": resolved}) +} + +// parseRawApprovalResponse parses the `response` field of an approval callback, +// which the control plane may deliver as either a JSON object or a JSON-encoded +// string. A non-JSON string is surfaced under a "text" key rather than dropped. +// Mirrors the TypeScript SDK's parseRawResponse and the Python SDK's handling. +func parseRawApprovalResponse(raw json.RawMessage) map[string]any { + if len(raw) == 0 { + return nil + } + // Direct object (or JSON null → nil). + var obj map[string]any + if err := json.Unmarshal(raw, &obj); err == nil { + return obj + } + // A JSON-encoded string that may itself contain a JSON object. + var s string + if err := json.Unmarshal(raw, &s); err == nil { + if strings.TrimSpace(s) == "" { + return nil + } + var inner map[string]any + if err := json.Unmarshal([]byte(s), &inner); err == nil { + return inner + } + return map[string]any{"text": s} + } + return nil +} diff --git a/sdk/go/agent/pause_test.go b/sdk/go/agent/pause_test.go new file mode 100644 index 000000000..d5f9aa64f --- /dev/null +++ b/sdk/go/agent/pause_test.go @@ -0,0 +1,637 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "io" + "log" + "net/http" + "net/http/httptest" + "strconv" + "sync" + "testing" + "time" + + "github.com/Agent-Field/agentfield/sdk/go/client" +) + +// newPauseAgent constructs a minimal Agent wired only with a PauseManager, +// mirroring newCancelAgent in cancel_test.go. Avoids the full New(cfg) +// bootstrap which drags in memory backends, AI clients, and DID subsystems we +// don't need here. +func newPauseAgent() *Agent { + return &Agent{pauseManager: NewPauseManager()} +} + +// --------------------------------------------------------------------------- +// ApprovalResult +// --------------------------------------------------------------------------- + +func TestApprovalResult_ConvenienceHelpers(t *testing.T) { + cases := []struct { + decision string + approved bool + changesRequested bool + }{ + {ApprovalApproved, true, false}, + {ApprovalRejected, false, false}, + {ApprovalRequestChanges, false, true}, + {ApprovalExpired, false, false}, + {ApprovalError, false, false}, + {ApprovalCancelled, false, false}, + } + for _, tc := range cases { + r := ApprovalResult{Decision: tc.decision} + if r.Approved() != tc.approved { + t.Errorf("decision %q: Approved()=%v want %v", tc.decision, r.Approved(), tc.approved) + } + if r.ChangesRequested() != tc.changesRequested { + t.Errorf("decision %q: ChangesRequested()=%v want %v", tc.decision, r.ChangesRequested(), tc.changesRequested) + } + } +} + +// --------------------------------------------------------------------------- +// PauseManager +// --------------------------------------------------------------------------- + +func TestPauseManager_ResolveByApprovalRequestID(t *testing.T) { + m := NewPauseManager() + ch := m.Register("req-1", "exec-1") + if got := m.PendingCount(); got != 1 { + t.Fatalf("PendingCount=%d want 1", got) + } + + if !m.Resolve("req-1", ApprovalResult{Decision: ApprovalApproved, Feedback: "lgtm"}) { + t.Fatal("Resolve returned false for a registered pause") + } + + result := <-ch + if !result.Approved() || result.Feedback != "lgtm" { + t.Fatalf("result=%+v want approved/lgtm", result) + } + if got := m.PendingCount(); got != 0 { + t.Fatalf("PendingCount=%d want 0 after resolve", got) + } +} + +func TestPauseManager_ResolveByExecutionID_Fallback(t *testing.T) { + m := NewPauseManager() + ch := m.Register("req-2", "exec-2") + + if !m.ResolveByExecutionID("exec-2", ApprovalResult{Decision: ApprovalRejected}) { + t.Fatal("ResolveByExecutionID returned false for a registered pause") + } + if got := (<-ch).Decision; got != ApprovalRejected { + t.Fatalf("decision=%q want rejected", got) + } +} + +func TestPauseManager_RegisterIsIdempotent(t *testing.T) { + m := NewPauseManager() + a := m.Register("req-3", "exec-3") + b := m.Register("req-3", "exec-3") + if m.PendingCount() != 1 { + t.Fatalf("PendingCount=%d want 1 after duplicate register", m.PendingCount()) + } + // Both channels must be the same underlying channel: resolving once + // delivers to a single receiver. + m.Resolve("req-3", ApprovalResult{Decision: ApprovalApproved}) + if got := (<-a).Decision; got != ApprovalApproved { + t.Fatalf("channel a decision=%q want approved", got) + } + select { + case _, ok := <-b: + if ok { + t.Fatal("second channel should be the same (closed, no second value)") + } + default: + t.Fatal("expected b to be the same closed channel as a") + } +} + +func TestPauseManager_ResolveUnknownReturnsFalse(t *testing.T) { + m := NewPauseManager() + if m.Resolve("nope", ApprovalResult{Decision: ApprovalApproved}) { + t.Fatal("Resolve returned true for unknown request") + } + if m.ResolveByExecutionID("nope", ApprovalResult{Decision: ApprovalApproved}) { + t.Fatal("ResolveByExecutionID returned true for unknown execution") + } +} + +func TestPauseManager_DoubleResolveReturnsFalse(t *testing.T) { + m := NewPauseManager() + ch := m.Register("req-dbl", "exec-dbl") + if !m.Resolve("req-dbl", ApprovalResult{Decision: ApprovalApproved}) { + t.Fatal("first Resolve returned false") + } + <-ch + if m.Resolve("req-dbl", ApprovalResult{Decision: ApprovalRejected}) { + t.Fatal("second Resolve returned true — must be idempotent") + } +} + +func TestPauseManager_CancelAll(t *testing.T) { + m := NewPauseManager() + c1 := m.Register("req-a", "exec-a") + c2 := m.Register("req-b", "exec-b") + + m.CancelAll() + if m.PendingCount() != 0 { + t.Fatalf("PendingCount=%d want 0 after CancelAll", m.PendingCount()) + } + if got := (<-c1).Decision; got != ApprovalCancelled { + t.Fatalf("c1 decision=%q want cancelled", got) + } + if got := (<-c2).Decision; got != ApprovalCancelled { + t.Fatalf("c2 decision=%q want cancelled", got) + } +} + +// Concurrent pauses must resolve independently — each waiter receives exactly +// its own result, with no cross-talk between approval_request_ids. +func TestPauseManager_ConcurrentPausesResolveIndependently(t *testing.T) { + m := NewPauseManager() + const n = 50 + + chans := make([]<-chan ApprovalResult, n) + for i := 0; i < n; i++ { + chans[i] = m.Register("req-"+strconv.Itoa(i), "exec-"+strconv.Itoa(i)) + } + if m.PendingCount() != n { + t.Fatalf("PendingCount=%d want %d", m.PendingCount(), n) + } + + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + m.Resolve("req-"+strconv.Itoa(i), ApprovalResult{ + Decision: ApprovalApproved, + Feedback: "fb-" + strconv.Itoa(i), + ApprovalRequestID: "req-" + strconv.Itoa(i), + }) + }(i) + } + wg.Wait() + + for i := 0; i < n; i++ { + got := <-chans[i] + want := "fb-" + strconv.Itoa(i) + if got.Feedback != want { + t.Fatalf("waiter %d received feedback %q, want %q (cross-talk)", i, got.Feedback, want) + } + } + if m.PendingCount() != 0 { + t.Fatalf("PendingCount=%d want 0 after all resolved", m.PendingCount()) + } +} + +// --------------------------------------------------------------------------- +// handleApprovalWebhook +// --------------------------------------------------------------------------- + +func postApprovalWebhook(t *testing.T, h http.HandlerFunc, body map[string]any) (*httptest.ResponseRecorder, map[string]any) { + t.Helper() + raw, err := json.Marshal(body) + if err != nil { + t.Fatalf("marshal body: %v", err) + } + req := httptest.NewRequest(http.MethodPost, "/webhooks/approval", bytes.NewReader(raw)) + rr := httptest.NewRecorder() + h(rr, req) + return rr, decodeJSON(t, rr.Body) +} + +func TestApprovalWebhook_ResolvesAndReplies(t *testing.T) { + a := newPauseAgent() + pending := a.pauseManager.Register("req-webhook", "exec-webhook") + + rr, resp := postApprovalWebhook(t, a.handleApprovalWebhook, map[string]any{ + "execution_id": "exec-webhook", + "approval_request_id": "req-webhook", + "decision": "approved", + "feedback": "ship it", + "response": `{"reviewer":"alice"}`, + }) + + if rr.Code != http.StatusOK { + t.Fatalf("status=%d want 200", rr.Code) + } + if resp["status"] != "received" || resp["resolved"] != true { + t.Fatalf("reply=%v want {status:received, resolved:true}", resp) + } + + result := <-pending + if !result.Approved() { + t.Errorf("result not approved: %+v", result) + } + if result.Feedback != "ship it" { + t.Errorf("feedback=%q want 'ship it'", result.Feedback) + } + if result.RawResponse["reviewer"] != "alice" { + t.Errorf("rawResponse=%v want reviewer=alice", result.RawResponse) + } +} + +// Decision strings map straight through to ApprovalResult.Decision. +func TestApprovalWebhook_DecisionMapping(t *testing.T) { + for _, decision := range []string{ApprovalApproved, ApprovalRejected, ApprovalRequestChanges, ApprovalExpired} { + t.Run(decision, func(t *testing.T) { + a := newPauseAgent() + pending := a.pauseManager.Register("req-"+decision, "exec-"+decision) + _, resp := postApprovalWebhook(t, a.handleApprovalWebhook, map[string]any{ + "execution_id": "exec-" + decision, + "approval_request_id": "req-" + decision, + "decision": decision, + }) + if resp["resolved"] != true { + t.Fatalf("resolved=%v want true", resp["resolved"]) + } + if got := (<-pending).Decision; got != decision { + t.Fatalf("decision=%q want %q", got, decision) + } + }) + } +} + +func TestApprovalWebhook_FallbackByExecutionID(t *testing.T) { + a := newPauseAgent() + pending := a.pauseManager.Register("req-fallback", "exec-fallback") + + // No approval_request_id in the callback — must fall back to execution_id. + _, resp := postApprovalWebhook(t, a.handleApprovalWebhook, map[string]any{ + "execution_id": "exec-fallback", + "decision": "request_changes", + }) + if resp["resolved"] != true { + t.Fatalf("resolved=%v want true", resp["resolved"]) + } + if !(<-pending).ChangesRequested() { + t.Fatal("expected changes-requested result") + } +} + +func TestApprovalWebhook_ResolvedFalseWhenNoMatch(t *testing.T) { + a := newPauseAgent() + _, resp := postApprovalWebhook(t, a.handleApprovalWebhook, map[string]any{ + "execution_id": "unknown", + "approval_request_id": "unknown", + "decision": "approved", + }) + if resp["resolved"] != false { + t.Fatalf("resolved=%v want false", resp["resolved"]) + } +} + +// The `response` field may be a JSON-encoded string or an inline object; both +// must land in RawResponse. A non-JSON string is surfaced under "text". +func TestApprovalWebhook_RawResponseParsing(t *testing.T) { + cases := []struct { + name string + response any + check func(t *testing.T, raw map[string]any) + }{ + { + name: "json-encoded-string", + response: `{"note":"stringy"}`, + check: func(t *testing.T, raw map[string]any) { + if raw["note"] != "stringy" { + t.Fatalf("raw=%v want note=stringy", raw) + } + }, + }, + { + name: "inline-object", + response: map[string]any{"note": "inline object"}, + check: func(t *testing.T, raw map[string]any) { + if raw["note"] != "inline object" { + t.Fatalf("raw=%v want note='inline object'", raw) + } + }, + }, + { + name: "non-json-string", + response: "just words", + check: func(t *testing.T, raw map[string]any) { + if raw["text"] != "just words" { + t.Fatalf("raw=%v want text='just words'", raw) + } + }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + a := newPauseAgent() + pending := a.pauseManager.Register("req-"+tc.name, "exec-"+tc.name) + postApprovalWebhook(t, a.handleApprovalWebhook, map[string]any{ + "execution_id": "exec-" + tc.name, + "approval_request_id": "req-" + tc.name, + "decision": "approved", + "response": tc.response, + }) + tc.check(t, (<-pending).RawResponse) + }) + } +} + +func TestApprovalWebhook_MissingFieldsReturns400(t *testing.T) { + a := newPauseAgent() + cases := []map[string]any{ + {"decision": "approved"}, // no execution_id + {"execution_id": "exec-x"}, // no decision + {}, // neither + } + for _, body := range cases { + rr, _ := postApprovalWebhook(t, a.handleApprovalWebhook, body) + if rr.Code != http.StatusBadRequest { + t.Fatalf("body %v: status=%d want 400", body, rr.Code) + } + } +} + +func TestApprovalWebhook_RejectsGet(t *testing.T) { + a := newPauseAgent() + req := httptest.NewRequest(http.MethodGet, "/webhooks/approval", nil) + rr := httptest.NewRecorder() + a.handleApprovalWebhook(rr, req) + if rr.Code != http.StatusMethodNotAllowed { + t.Fatalf("status=%d want 405", rr.Code) + } +} + +// The approval webhook route is always wired into the mux (independent of any +// accepts_webhook trigger flag), matching the Python/TS SDKs where it is +// auto-installed for every agent. A POST with missing fields returns 400 (route +// present) rather than 404 (route absent). +func TestApprovalWebhook_RouteAlwaysRegisteredOnMux(t *testing.T) { + a := newPauseAgent() + srv := httptest.NewServer(a.Handler()) + defer srv.Close() + + resp, err := http.Post(srv.URL+"/webhooks/approval", "application/json", + bytes.NewReader([]byte(`{"decision":"approved"}`))) + if err != nil { + t.Fatalf("POST failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("status=%d want 400 (route present but missing fields)", resp.StatusCode) + } +} + +// The control plane delivers the approval callback unauthenticated, so the +// /webhooks/approval route must bypass origin-token auth — otherwise the +// callback is rejected with 401 and the pause never resolves. +func TestApprovalWebhook_BypassesOriginAuth(t *testing.T) { + a := newPauseAgent() + a.cfg.RequireOriginAuth = true + a.cfg.InternalToken = "secret-token" + srv := httptest.NewServer(a.Handler()) + defer srv.Close() + + // No Authorization header — an ordinary /execute call would 401 here. + resp, err := http.Post(srv.URL+"/webhooks/approval", "application/json", + bytes.NewReader([]byte(`{"execution_id":"e","approval_request_id":"r","decision":"approved"}`))) + if err != nil { + t.Fatalf("POST failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized { + t.Fatal("approval webhook was rejected by origin auth — must be exempt") + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d want 200 (handler reached)", resp.StatusCode) + } +} + +// Likewise the route must bypass local DID verification — the CP callback +// carries no DID signature. +func TestApprovalWebhook_BypassesLocalVerification(t *testing.T) { + a := newPauseAgent() + // A verifier with no reachable control plane; the exemption must short + // circuit before any DID/refresh logic runs, so no network call happens. + a.localVerifier = NewLocalVerifier("http://127.0.0.1:0", time.Hour, "") + a.logger = log.New(io.Discard, "", 0) + srv := httptest.NewServer(a.Handler()) + defer srv.Close() + + resp, err := http.Post(srv.URL+"/webhooks/approval", "application/json", + bytes.NewReader([]byte(`{"execution_id":"e","approval_request_id":"r","decision":"approved"}`))) + if err != nil { + t.Fatalf("POST failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + t.Fatalf("approval webhook blocked by local verification (status=%d) — must be exempt", resp.StatusCode) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status=%d want 200 (handler reached)", resp.StatusCode) + } +} + +// --------------------------------------------------------------------------- +// Agent.Pause — end-to-end through the control-plane request-approval call and +// the webhook callback. +// --------------------------------------------------------------------------- + +// mockControlPlane returns an httptest server that answers request-approval. +// If fireCallback is true it POSTs the given resolution to the callback_url it +// received, exercising the full pause/resolve loop. +func mockControlPlane(t *testing.T, status int, fireCallback bool, resolution map[string]any) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req client.RequestApprovalRequest + _ = json.NewDecoder(r.Body).Decode(&req) + if status >= 400 { + w.WriteHeader(status) + _, _ = w.Write([]byte(`{"error":"boom"}`)) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]any{ + "approval_request_id": req.ApprovalRequestID, + "approval_request_url": req.ApprovalRequestURL, + "status": "waiting", + }) + if fireCallback && req.CallbackURL != "" { + payload := map[string]any{"execution_id": "exec-e2e", "approval_request_id": req.ApprovalRequestID} + for k, v := range resolution { + payload[k] = v + } + body, _ := json.Marshal(payload) + go func() { + // Best-effort async callback, like the real control plane. + _, _ = http.Post(req.CallbackURL, "application/json", bytes.NewReader(body)) + }() + } + })) +} + +// buildPauseAgent wires a minimal Agent with a control-plane client and an +// HTTP server for the /webhooks/approval callback. Returns the agent and a +// cleanup func. +func buildPauseAgent(t *testing.T, cpURL string) (*Agent, func()) { + t.Helper() + c, err := client.New(cpURL) + if err != nil { + t.Fatalf("client.New: %v", err) + } + a := &Agent{ + pauseManager: NewPauseManager(), + client: c, + } + a.cfg.NodeID = "test-node" + srv := httptest.NewServer(a.Handler()) + a.cfg.PublicURL = srv.URL + return a, srv.Close +} + +func TestPause_ResolvesOnWebhookCallback(t *testing.T) { + cp := mockControlPlane(t, http.StatusOK, true, map[string]any{ + "decision": "approved", + "feedback": "lgtm", + "response": `{"reviewer":"bob"}`, + }) + defer cp.Close() + + a, cleanup := buildPauseAgent(t, cp.URL) + defer cleanup() + + result, err := a.Pause(context.Background(), PauseOptions{ + ApprovalRequestID: "req-e2e", + ExecutionID: "exec-e2e", + Timeout: 5 * time.Second, + }) + if err != nil { + t.Fatalf("Pause returned error: %v", err) + } + if !result.Approved() { + t.Fatalf("result=%+v want approved", result) + } + if result.Feedback != "lgtm" { + t.Errorf("feedback=%q want lgtm", result.Feedback) + } + if result.RawResponse["reviewer"] != "bob" { + t.Errorf("rawResponse=%v want reviewer=bob", result.RawResponse) + } + if a.pauseManager.PendingCount() != 0 { + t.Errorf("PendingCount=%d want 0", a.pauseManager.PendingCount()) + } +} + +func TestPause_ExpiryReturnsExpiredDecision(t *testing.T) { + // CP accepts request-approval but never fires a callback. + cp := mockControlPlane(t, http.StatusOK, false, nil) + defer cp.Close() + + a, cleanup := buildPauseAgent(t, cp.URL) + defer cleanup() + + result, err := a.Pause(context.Background(), PauseOptions{ + ApprovalRequestID: "req-expire", + ExecutionID: "exec-e2e", + Timeout: 50 * time.Millisecond, + }) + if err != nil { + t.Fatalf("Pause returned error on timeout, want nil: %v", err) + } + if result.Decision != ApprovalExpired { + t.Fatalf("decision=%q want expired", result.Decision) + } + if a.pauseManager.PendingCount() != 0 { + t.Errorf("PendingCount=%d want 0 after expiry", a.pauseManager.PendingCount()) + } +} + +func TestPause_ContextCancellation(t *testing.T) { + cp := mockControlPlane(t, http.StatusOK, false, nil) + defer cp.Close() + + a, cleanup := buildPauseAgent(t, cp.URL) + defer cleanup() + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + result, err := a.Pause(ctx, PauseOptions{ + ApprovalRequestID: "req-cancel", + ExecutionID: "exec-e2e", + Timeout: 10 * time.Second, + }) + if err == nil { + t.Fatalf("Pause returned nil error on cancellation, result=%+v", result) + } + if err != context.Canceled { + t.Fatalf("err=%v want context.Canceled", err) + } + if a.pauseManager.PendingCount() != 0 { + t.Errorf("PendingCount=%d want 0 after cancellation", a.pauseManager.PendingCount()) + } +} + +func TestPause_RequestApprovalFailureReturnsError(t *testing.T) { + cp := mockControlPlane(t, http.StatusInternalServerError, false, nil) + defer cp.Close() + + a, cleanup := buildPauseAgent(t, cp.URL) + defer cleanup() + + _, err := a.Pause(context.Background(), PauseOptions{ + ApprovalRequestID: "req-fail", + ExecutionID: "exec-e2e", + Timeout: time.Second, + }) + if err == nil { + t.Fatal("Pause returned nil error when control plane failed") + } + // The pending pause must have been cleaned up. + if a.pauseManager.PendingCount() != 0 { + t.Errorf("PendingCount=%d want 0 after CP failure", a.pauseManager.PendingCount()) + } +} + +func TestPause_NoExecutionIDReturnsError(t *testing.T) { + a := newPauseAgent() + _, err := a.Pause(context.Background(), PauseOptions{ApprovalRequestID: "req-x"}) + if err == nil { + t.Fatal("expected error when no execution_id is available") + } +} + +func TestPause_MissingApprovalRequestIDReturnsError(t *testing.T) { + a := newPauseAgent() + _, err := a.Pause(context.Background(), PauseOptions{ExecutionID: "exec-1"}) + if err == nil { + t.Fatal("expected error when approval_request_id is empty") + } +} + +// Pause resolves the execution_id from the ExecutionContext on ctx when +// PauseOptions.ExecutionID is not set. +func TestPause_UsesExecutionIDFromContext(t *testing.T) { + cp := mockControlPlane(t, http.StatusOK, true, map[string]any{"decision": "approved"}) + defer cp.Close() + + a, cleanup := buildPauseAgent(t, cp.URL) + defer cleanup() + + ctx := contextWithExecution(context.Background(), ExecutionContext{ExecutionID: "exec-e2e"}) + result, err := a.Pause(ctx, PauseOptions{ + ApprovalRequestID: "req-ctx", + Timeout: 5 * time.Second, + }) + if err != nil { + t.Fatalf("Pause error: %v", err) + } + if !result.Approved() { + t.Fatalf("result=%+v want approved", result) + } +} From efb0aad552b8e71df05cb3e974c6e7876ff5afc0 Mon Sep 17 00:00:00 2001 From: Abir Abbas Date: Thu, 9 Jul 2026 16:57:25 -0400 Subject: [PATCH 2/2] feat(sdk/go): harness cost reporting and incremental schema mode Metrics/Result gain CostUSD (*float64, nil = unknown) extracted from Claude's JSON output with Python's cost_usd-or-total_cost_usd semantics and accumulated across retries including failed attempts. Options gains SchemaMode (single/incremental/auto): auto engages on the compact-encoded schema crossing the large-schema token threshold, with the incremental prompt suffix, per-field failure diagnosis, and followup prompts ported byte-verbatim from the Python SDK. Non-Claude providers report nil cost (Python parity when litellm is unavailable). Co-Authored-By: Claude Fable 5 --- sdk/go/harness/claudecode.go | 29 ++ sdk/go/harness/coverage_branches_test.go | 35 +-- sdk/go/harness/coverage_helpers_test.go | 4 +- sdk/go/harness/parity_test.go | 355 +++++++++++++++++++++++ sdk/go/harness/provider.go | 13 + sdk/go/harness/result.go | 15 + sdk/go/harness/runner.go | 99 ++++++- sdk/go/harness/runner_invariant_test.go | 12 +- sdk/go/harness/runner_test.go | 8 +- sdk/go/harness/schema.go | 222 ++++++++++++++ 10 files changed, 756 insertions(+), 36 deletions(-) create mode 100644 sdk/go/harness/parity_test.go diff --git a/sdk/go/harness/claudecode.go b/sdk/go/harness/claudecode.go index 7a0ece23d..debd5f96b 100644 --- a/sdk/go/harness/claudecode.go +++ b/sdk/go/harness/claudecode.go @@ -150,6 +150,7 @@ func (p *ClaudeCodeProvider) parseJSONOutput(stdout string, raw *RawResult) { var messages []map[string]any var resultText string var sessionID string + var cost *float64 numTurns := 0 for _, line := range strings.Split(stdout, "\n") { @@ -174,6 +175,9 @@ func (p *ClaudeCodeProvider) parseJSONOutput(stdout string, raw *RawResult) { if sid, ok := msg["session_id"].(string); ok { sessionID = sid } + if c := extractCost(msg); c != nil { + cost = c + } if turns, ok := msg["num_turns"].(float64); ok { numTurns = int(turns) } @@ -189,12 +193,37 @@ func (p *ClaudeCodeProvider) parseJSONOutput(stdout string, raw *RawResult) { } raw.Messages = messages raw.Metrics.SessionID = sessionID + raw.Metrics.CostUSD = cost raw.Metrics.NumTurns = numTurns if numTurns == 0 && len(messages) > 0 { raw.Metrics.NumTurns = len(messages) } } +// extractCost pulls the per-call cost from a Claude Code result message. +// +// Mirrors the Python provider's semantics exactly +// (agentfield/harness/providers/claude.py): +// +// cost_info = msg.get("cost_usd") or msg.get("total_cost_usd") +// if cost_info is not None: +// total_cost = float(cost_info) +// +// The Python `or` treats a zero/absent "cost_usd" as falsy and falls through +// to "total_cost_usd". Returns nil when neither yields a usable number, so the +// caller can distinguish "unknown cost" (nil) from "$0.00". +func extractCost(msg map[string]any) *float64 { + if v, ok := msg["cost_usd"].(float64); ok && v != 0 { + c := v + return &c + } + if v, ok := msg["total_cost_usd"].(float64); ok { + c := v + return &c + } + return nil +} + // extractAssistantText pulls text content from an assistant message. func extractAssistantText(msg map[string]any) string { // Direct content field diff --git a/sdk/go/harness/coverage_branches_test.go b/sdk/go/harness/coverage_branches_test.go index 7770c1a5f..043e1e4c5 100644 --- a/sdk/go/harness/coverage_branches_test.go +++ b/sdk/go/harness/coverage_branches_test.go @@ -18,23 +18,23 @@ func TestResultText(t *testing.T) { func TestRunnerMergeOptions_AllOverrideBranches(t *testing.T) { defaults := Options{ - Provider: ProviderOpenCode, - Model: "base-model", - MaxTurns: 1, - PermissionMode: "plan", - SystemPrompt: "base system", - Env: map[string]string{"BASE": "1"}, - Cwd: "/base/cwd", - ProjectDir: "/base/project", - Tools: []string{"base-tool"}, - MaxBudgetUSD: 1.25, - ResumeSessionID: "base-session", - BinPath: "/base/bin", - Timeout: 10, - MaxRetries: 1, - InitialDelay: 1.5, - MaxDelay: 2.5, - BackoffFactor: 3.5, + Provider: ProviderOpenCode, + Model: "base-model", + MaxTurns: 1, + PermissionMode: "plan", + SystemPrompt: "base system", + Env: map[string]string{"BASE": "1"}, + Cwd: "/base/cwd", + ProjectDir: "/base/project", + Tools: []string{"base-tool"}, + MaxBudgetUSD: 1.25, + ResumeSessionID: "base-session", + BinPath: "/base/bin", + Timeout: 10, + MaxRetries: 1, + InitialDelay: 1.5, + MaxDelay: 2.5, + BackoffFactor: 3.5, SchemaMaxRetries: 1, } @@ -296,6 +296,7 @@ func TestRunnerRetryAdditionalBranches(t *testing.T) { prov, Options{SchemaMaxRetries: 2}, "prompt", + false, ) assert.True(t, result.IsError) diff --git a/sdk/go/harness/coverage_helpers_test.go b/sdk/go/harness/coverage_helpers_test.go index 596cc8621..13c0a54fc 100644 --- a/sdk/go/harness/coverage_helpers_test.go +++ b/sdk/go/harness/coverage_helpers_test.go @@ -217,10 +217,11 @@ func TestSchemaHelperBranches(t *testing.T) { func TestRunnerHelperBranches(t *testing.T) { t.Run("accumulateMetrics merges turns, session ids, and messages", func(t *testing.T) { - turns, sid, msgs := accumulateMetrics([]*RawResult{ + cost, turns, sid, msgs := accumulateMetrics([]*RawResult{ {Metrics: Metrics{NumTurns: 1, SessionID: "old"}, Messages: []map[string]any{{"a": 1}}}, {Metrics: Metrics{NumTurns: 2, SessionID: "new"}, Messages: []map[string]any{{"b": 2}}}, }) + assert.Nil(t, cost) // no execution reported a cost assert.Equal(t, 3, turns) assert.Equal(t, "new", sid) assert.Len(t, msgs, 2) @@ -311,6 +312,7 @@ func TestRunnerHelperBranches(t *testing.T) { }}, Options{SchemaMaxRetries: 2}, "prompt", + false, ) assert.True(t, result.IsError) diff --git a/sdk/go/harness/parity_test.go b/sdk/go/harness/parity_test.go new file mode 100644 index 000000000..b9df3715c --- /dev/null +++ b/sdk/go/harness/parity_test.go @@ -0,0 +1,355 @@ +package harness + +import ( + "context" + "encoding/json" + "fmt" + "os" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- Per-call cost reporting --------------------------------------------- + +// TestExtractCost pins the Python `cost_usd or total_cost_usd` semantics from +// agentfield/harness/providers/claude.py. +func TestExtractCost(t *testing.T) { + t.Run("total_cost_usd only", func(t *testing.T) { + c := extractCost(map[string]any{"total_cost_usd": 0.42}) + require.NotNil(t, c) + assert.InDelta(t, 0.42, *c, 1e-9) + }) + + t.Run("cost_usd takes precedence when non-zero", func(t *testing.T) { + c := extractCost(map[string]any{"cost_usd": 1.5, "total_cost_usd": 9.9}) + require.NotNil(t, c) + assert.InDelta(t, 1.5, *c, 1e-9) + }) + + t.Run("zero cost_usd falls through to total_cost_usd", func(t *testing.T) { + c := extractCost(map[string]any{"cost_usd": 0.0, "total_cost_usd": 2.0}) + require.NotNil(t, c) + assert.InDelta(t, 2.0, *c, 1e-9) + }) + + t.Run("neither present yields nil", func(t *testing.T) { + assert.Nil(t, extractCost(map[string]any{"type": "result"})) + }) + + t.Run("zero cost_usd with no total yields nil (0.0 or None == None)", func(t *testing.T) { + assert.Nil(t, extractCost(map[string]any{"cost_usd": 0.0})) + }) +} + +// TestClaudeCodeProvider_ExtractsCost verifies cost is pulled out of a Claude +// Code JSON result fixture (contract: cost extracted from a claude JSON output). +func TestClaudeCodeProvider_ExtractsCost(t *testing.T) { + dir := t.TempDir() + script := writeTestScript(t, dir, "claude", + `#!/bin/sh +echo '{"type":"result","result":"ok","session_id":"s1","num_turns":2,"total_cost_usd":0.0731}' +`) + + p := NewClaudeCodeProvider(script) + raw, err := p.Execute(context.Background(), "prompt", Options{}) + require.NoError(t, err) + require.False(t, raw.IsError) + require.NotNil(t, raw.Metrics.CostUSD) + assert.InDelta(t, 0.0731, *raw.Metrics.CostUSD, 1e-9) +} + +// TestClaudeCodeProvider_NoCostIsNil confirms nil (unknown), not 0.0, when the +// provider does not report cost. +func TestClaudeCodeProvider_NoCostIsNil(t *testing.T) { + dir := t.TempDir() + script := writeTestScript(t, dir, "claude", + `#!/bin/sh +echo '{"type":"result","result":"ok","session_id":"s1","num_turns":1}' +`) + + p := NewClaudeCodeProvider(script) + raw, err := p.Execute(context.Background(), "prompt", Options{}) + require.NoError(t, err) + require.False(t, raw.IsError) + assert.Nil(t, raw.Metrics.CostUSD) +} + +// TestAccumulateMetrics_Cost verifies cost sums over attempts and stays nil +// when nothing reported a cost (Python _accumulate_metrics semantics). +func TestAccumulateMetrics_Cost(t *testing.T) { + t.Run("nil when no attempt reports cost", func(t *testing.T) { + cost, _, _, _ := accumulateMetrics([]*RawResult{ + {Metrics: Metrics{NumTurns: 1}}, + {Metrics: Metrics{NumTurns: 1}}, + }) + assert.Nil(t, cost) + }) + + t.Run("sums present costs", func(t *testing.T) { + c1, c2 := 0.10, 0.25 + cost, _, _, _ := accumulateMetrics([]*RawResult{ + {Metrics: Metrics{CostUSD: &c1}}, + {Metrics: Metrics{NumTurns: 1}}, // no cost — skipped + {Metrics: Metrics{CostUSD: &c2}}, + }) + require.NotNil(t, cost) + assert.InDelta(t, 0.35, *cost, 1e-9) + }) +} + +// TestHandleSchemaWithRetry_AccumulatesCostIncludingFailedAttempts proves the +// end-to-end result cost sums across every attempt, INCLUDING failed retry +// attempts (Python appends retry_raw to all_raws before the is_error check). +func TestHandleSchemaWithRetry_AccumulatesCostIncludingFailedAttempts(t *testing.T) { + dir := t.TempDir() + type Out struct { + Value string `json:"value"` + } + schema := map[string]any{ + "type": "object", + "properties": map[string]any{"value": map[string]any{"type": "string"}}, + } + + c0, c1, c2 := 0.10, 0.20, 0.30 + initialRaw := &RawResult{Result: "no json", Metrics: Metrics{NumTurns: 1, CostUSD: &c0}} + mock := &mockProvider{results: []*RawResult{ + {Result: "still bad", IsError: true, ErrorMessage: "nope", Metrics: Metrics{NumTurns: 1, CostUSD: &c1}}, + {Result: "still bad", IsError: true, ErrorMessage: "nope", Metrics: Metrics{NumTurns: 1, CostUSD: &c2}}, + }} + + var dest Out + result := NewRunner(Options{Provider: "opencode"}).handleSchemaWithRetry( + context.Background(), initialRaw, schema, &dest, dir, + time.Now(), mock, Options{Provider: "opencode", SchemaMaxRetries: 2}, "test prompt", false, + ) + + require.True(t, result.IsError) + require.NotNil(t, result.CostUSD) + // 0.10 (initial) + 0.20 + 0.30 (both failed retries) = 0.60 + assert.InDelta(t, 0.60, *result.CostUSD, 1e-9) +} + +// --- Incremental schema mode --------------------------------------------- + +// TestResolveIncremental covers the mode-selection rules from the Python +// HarnessRunner._resolve_incremental. +func TestResolveIncremental(t *testing.T) { + small := map[string]any{ + "type": "object", + "properties": map[string]any{"a": map[string]any{"type": "string"}}, + } + + assert.False(t, resolveIncremental(nil, Options{SchemaMode: "incremental"}), "nil schema never incremental") + assert.True(t, resolveIncremental(small, Options{SchemaMode: "incremental"})) + assert.True(t, resolveIncremental(small, Options{SchemaMode: "INCREMENTAL"}), "mode is case-insensitive") + assert.False(t, resolveIncremental(small, Options{}), "default (empty) is single-shot") + assert.False(t, resolveIncremental(small, Options{SchemaMode: "single"})) + assert.False(t, resolveIncremental(small, Options{SchemaMode: "auto"}), "small schema stays single-shot under auto") +} + +// TestResolveIncremental_AutoEngagesOnLargeSchema constructs the auto trigger: +// a schema whose compact JSON exceeds the large-schema token threshold engages +// incremental, matching Python (which measures the compact json.dumps output). +func TestResolveIncremental_AutoEngagesOnLargeSchema(t *testing.T) { + large := largeParitySchema() + + compact, err := json.Marshal(large) + require.NoError(t, err) + require.Greater(t, estimateTokens(string(compact)), largeSchemaTokenThreshold, + "fixture must exceed the threshold on its COMPACT encoding") + + assert.True(t, resolveIncremental(large, Options{SchemaMode: "auto"})) +} + +// TestBuildIncrementalPromptSuffix checks the byte-verbatim incremental build +// instructions and the per-field required/optional listing. +func TestBuildIncrementalPromptSuffix(t *testing.T) { + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "title": map[string]any{"type": "string"}, + "body": map[string]any{"type": "string"}, + }, + "required": []any{"title"}, + } + dir := t.TempDir() + s := BuildIncrementalPromptSuffix(schema, dir) + + assert.Contains(t, s, "CRITICAL OUTPUT REQUIREMENTS (incremental build):") + assert.Contains(t, s, "Build it ONE FIELD AT A TIME so nothing gets truncated:") + assert.Contains(t, s, " 1. First create the file with an empty object: {}") + assert.Contains(t, s, " each edit re-read the file to confirm it is still valid JSON.") + assert.Contains(t, s, " 4. Do not finish until every required field is present.") + assert.Contains(t, s, "Top-level fields to add:") + assert.Contains(t, s, " - title (required)") + assert.Contains(t, s, " - body (optional)") + assert.Contains(t, s, "Full JSON Schema:") + assert.Contains(t, s, OutputPath(dir)) + assert.Contains(t, s, "no markdown fences, no commentary, no extra text.") +} + +// TestBuildIncrementalFollowup checks the byte-verbatim patch-only follow-up. +func TestBuildIncrementalFollowup(t *testing.T) { + schema := map[string]any{ + "type": "object", + "properties": map[string]any{"body": map[string]any{"type": "string"}}, + } + dir := t.TempDir() + s := BuildIncrementalFollowup(map[string]string{"body": "missing required field"}, dir, schema) + + assert.Contains(t, s, "PARTIAL OUTPUT NEEDS FIXES. The JSON at ") + assert.Contains(t, s, " is incomplete or invalid.") + assert.Contains(t, s, "Patch ONLY these fields, one at a time, using Edit, keeping the file valid JSON after each change:") + assert.Contains(t, s, " - body: missing required field") + assert.Contains(t, s, "Full schema:") + assert.Contains(t, s, "Leave every already-correct field unchanged. Do NOT rewrite the whole file.") +} + +// TestDiagnoseFieldFailures covers the field-level diagnosis reasons used to +// drive incremental recovery. +func TestDiagnoseFieldFailures(t *testing.T) { + dir := t.TempDir() + type Out struct { + Title string `json:"title"` + Body string `json:"body"` + } + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "title": map[string]any{"type": "string"}, + "body": map[string]any{"type": "string"}, + }, + "required": []any{"title", "body"}, + } + path := OutputPath(dir) + + t.Run("missing required field", func(t *testing.T) { + require.NoError(t, os.WriteFile(path, []byte(`{"title":"x"}`), 0o644)) + var dest Out + f := DiagnoseFieldFailures(path, schema, &dest) + assert.Equal(t, "missing required field", f["body"]) + _, hasTitle := f["title"] + assert.False(t, hasTitle, "present required field is not flagged") + }) + + t.Run("file not a JSON object reports required fields", func(t *testing.T) { + require.NoError(t, os.WriteFile(path, []byte(`not json at all`), 0o644)) + var dest Out + f := DiagnoseFieldFailures(path, schema, &dest) + assert.Equal(t, "output file missing or not a JSON object", f["title"]) + assert.Equal(t, "output file missing or not a JSON object", f["body"]) + }) + + t.Run("clean object has no failures", func(t *testing.T) { + require.NoError(t, os.WriteFile(path, []byte(`{"title":"x","body":"y"}`), 0o644)) + var dest Out + f := DiagnoseFieldFailures(path, schema, &dest) + assert.Empty(t, f) + }) + + t.Run("type mismatch reports _root", func(t *testing.T) { + require.NoError(t, os.WriteFile(path, []byte(`{"title":"x","body":123}`), 0o644)) + var dest Out + f := DiagnoseFieldFailures(path, schema, &dest) + _, has := f["_root"] + assert.True(t, has) + // The caller's dest must NOT be mutated by diagnosis. + assert.Equal(t, "", dest.Title) + }) +} + +// TestHandleSchemaWithRetry_IncrementalRecovery drives the incremental retry +// path: an initial partial/invalid output file triggers a patch-only follow-up +// (with the original goal preserved), and the corrected file assembles a valid +// dest. +func TestHandleSchemaWithRetry_IncrementalRecovery(t *testing.T) { + dir := t.TempDir() + type Out struct { + Title string `json:"title"` + Body string `json:"body"` + } + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "title": map[string]any{"type": "string"}, + "body": map[string]any{"type": "string"}, + }, + "required": []any{"title", "body"}, + } + outputPath := OutputPath(dir) + // Initial file has a type-mismatched required field, so validation fails + // and the incremental recovery path runs (Go struct unmarshal, unlike + // pydantic, does not fail on a merely-missing field). + require.NoError(t, os.WriteFile(outputPath, []byte(`{"title":"hello","body":123}`), 0o644)) + + var capturedPrompt string + prov := &funcProvider{fn: func(_ context.Context, prompt string, _ Options) (*RawResult, error) { + capturedPrompt = prompt + // Agent patches the file to a valid object. + _ = os.WriteFile(outputPath, []byte(`{"title":"hello","body":"world"}`), 0o644) + return &RawResult{Result: "patched", Metrics: Metrics{NumTurns: 1}}, nil + }} + + var dest Out + result := NewRunner(Options{}).handleSchemaWithRetry( + context.Background(), + &RawResult{Result: "", Metrics: Metrics{NumTurns: 1}}, + schema, &dest, dir, time.Now(), prov, + Options{SchemaMaxRetries: 2}, "ORIGINAL GOAL", true, // useIncremental + ) + + require.False(t, result.IsError) + assert.Equal(t, "hello", dest.Title) + assert.Equal(t, "world", dest.Body) + + assert.Contains(t, capturedPrompt, "PARTIAL OUTPUT NEEDS FIXES") + assert.Contains(t, capturedPrompt, "Patch ONLY these fields") + assert.Contains(t, capturedPrompt, "ORIGINAL GOAL", "original goal is prepended on incremental retry") +} + +// TestSchemaMode_SingleShotUnchanged asserts the single-shot production path is +// selected and unchanged when the mode is off. +func TestSchemaMode_SingleShotUnchanged(t *testing.T) { + schema := map[string]any{ + "type": "object", + "properties": map[string]any{"name": map[string]any{"type": "string"}}, + } + + assert.False(t, resolveIncremental(schema, Options{})) + assert.False(t, resolveIncremental(schema, Options{SchemaMode: "single"})) + + single := BuildPromptSuffix(schema, "/tmp/parity") + assert.Contains(t, single, "CRITICAL OUTPUT REQUIREMENTS:") + assert.NotContains(t, single, "incremental build") + + inc := BuildIncrementalPromptSuffix(schema, "/tmp/parity") + assert.Contains(t, inc, "CRITICAL OUTPUT REQUIREMENTS (incremental build):") +} + +// TestMergeOptions_SchemaMode verifies SchemaMode flows through the default/ +// override merge. +func TestMergeOptions_SchemaMode(t *testing.T) { + r := NewRunner(Options{Provider: "opencode", SchemaMode: "auto"}) + assert.Equal(t, "auto", r.mergeOptions(Options{}).SchemaMode, "default carried through") + assert.Equal(t, "incremental", r.mergeOptions(Options{SchemaMode: "incremental"}).SchemaMode, "override wins") +} + +// largeParitySchema builds a JSON schema whose COMPACT encoding exceeds the +// large-schema token threshold (~16000 chars). +func largeParitySchema() map[string]any { + props := make(map[string]any) + for i := 0; i < 600; i++ { + key := fmt.Sprintf("field_%04d", i) + props[key] = map[string]any{ + "type": "string", + "description": "A padded description that helps push the schema past the large-schema token threshold.", + } + } + return map[string]any{ + "type": "object", + "properties": props, + } +} diff --git a/sdk/go/harness/provider.go b/sdk/go/harness/provider.go index d5e1036ff..8e0a0dee9 100644 --- a/sdk/go/harness/provider.go +++ b/sdk/go/harness/provider.go @@ -80,6 +80,19 @@ type Options struct { // SchemaMaxRetries controls how many times to retry when schema // validation fails. Default 2. SchemaMaxRetries int + + // SchemaMode selects how schema-constrained output is produced: + // + // "single" (default, or "") — the agent writes the whole JSON object in + // one shot (cheapest). + // "incremental" — the agent builds the object one top-level + // field at a time, and recovery patches only the failing fields + // (more robust for large or deeply nested schemas). + // "auto" — incremental only when the schema is large + // (compact JSON token estimate exceeds the large-schema threshold). + // + // Mirrors the Python SDK's schema_mode argument to .harness(). + SchemaMode string } func (o Options) maxRetries() int { diff --git a/sdk/go/harness/result.go b/sdk/go/harness/result.go index 8b1762839..2d12cd914 100644 --- a/sdk/go/harness/result.go +++ b/sdk/go/harness/result.go @@ -21,6 +21,15 @@ type Metrics struct { DurationAPIMS int NumTurns int SessionID string + + // CostUSD is the cost in USD reported by the provider for this single + // execution, or nil when the provider does not report a cost. Mirrors the + // Python SDK's Metrics.total_cost_usd (Optional[float]): nil means + // "unknown", not "free". Only providers that surface a native cost (e.g. + // Claude Code's result JSON) populate this; providers whose Python + // counterparts derive cost via litellm token estimation leave it nil in Go + // because there is no Go equivalent of litellm's pricing database. + CostUSD *float64 } // RawResult is the output from a single provider execution before schema @@ -49,6 +58,12 @@ type Result struct { ErrorMessage string FailureType FailureType + // CostUSD is the total cost in USD accumulated across every provider + // execution that contributed to this result (including failed retry + // attempts), or nil when no execution reported a cost. Mirrors the Python + // SDK's HarnessResult.cost_usd (Optional[float]). + CostUSD *float64 + NumTurns int DurationMS int SessionID string diff --git a/sdk/go/harness/runner.go b/sdk/go/harness/runner.go index e0914dc2f..9afa3b322 100644 --- a/sdk/go/harness/runner.go +++ b/sdk/go/harness/runner.go @@ -2,6 +2,7 @@ package harness import ( "context" + "encoding/json" "fmt" "io" "log" @@ -65,9 +66,20 @@ func (r *Runner) Run(ctx context.Context, prompt string, schema map[string]any, outputDir = tempOutputDir } + // schema_mode selects how the agent is asked to produce the output: + // "single" — one Write of the whole object (default, cheapest) + // "incremental" — build the object one top-level field at a time, with + // field-level recovery (robust for large/deep schemas) + // "auto" — incremental only when the schema is large, else single + useIncremental := resolveIncremental(schema, opts) + effectivePrompt := prompt if schema != nil { - effectivePrompt = prompt + BuildPromptSuffix(schema, outputDir) + if useIncremental { + effectivePrompt = prompt + BuildIncrementalPromptSuffix(schema, outputDir) + } else { + effectivePrompt = prompt + BuildPromptSuffix(schema, outputDir) + } } startTime := time.Now() @@ -78,7 +90,7 @@ func (r *Runner) Run(ctx context.Context, prompt string, schema map[string]any, } if schema != nil { - result := r.handleSchemaWithRetry(ctx, raw, schema, dest, outputDir, startTime, provider, opts, effectivePrompt) + result := r.handleSchemaWithRetry(ctx, raw, schema, dest, outputDir, startTime, provider, opts, effectivePrompt, useIncremental) CleanupTempFiles(outputDir) return result, nil } @@ -89,6 +101,7 @@ func (r *Runner) Run(ctx context.Context, prompt string, schema map[string]any, IsError: raw.IsError, ErrorMessage: raw.ErrorMessage, FailureType: raw.FailureType, + CostUSD: raw.Metrics.CostUSD, NumTurns: raw.Metrics.NumTurns, DurationMS: elapsed, SessionID: raw.Metrics.SessionID, @@ -169,10 +182,48 @@ func (r *Runner) mergeOptions(overrides Options) Options { if overrides.SchemaMaxRetries > 0 { merged.SchemaMaxRetries = overrides.SchemaMaxRetries } + if overrides.SchemaMode != "" { + merged.SchemaMode = overrides.SchemaMode + } return merged } +// resolveIncremental decides whether to use the incremental field-by-field +// schema build. Mirrors the Python HarnessRunner._resolve_incremental: +// +// schema is nil -> false +// mode == "incremental" -> true +// mode == "auto" -> true iff the compact JSON schema exceeds the +// large-schema token threshold +// otherwise ("single"/"") -> false +// +// The "auto" branch uses the COMPACT JSON encoding (json.Marshal, matching +// Python's json.dumps with no indent), which differs from the indented +// encoding the prompt-suffix builders use to decide whether to spill the +// schema to a file — this matches Python exactly. +func resolveIncremental(jsonSchema map[string]any, opts Options) bool { + if jsonSchema == nil { + return false + } + mode := strings.ToLower(opts.SchemaMode) + if mode == "" { + mode = "single" + } + switch mode { + case "incremental": + return true + case "auto": + compact, err := json.Marshal(jsonSchema) + if err != nil { + return false + } + return estimateTokens(string(compact)) > largeSchemaTokenThreshold + default: + return false + } +} + // transientPatterns are substrings that indicate a retryable error. var transientPatterns = []string{ "rate limit", "rate_limit", "overloaded", "timeout", "timed out", @@ -252,6 +303,7 @@ func (r *Runner) handleSchemaWithRetry( provider Provider, opts Options, originalPrompt string, + useIncremental bool, ) *Result { outputPath := OutputPath(outputDir) maxRetries := opts.schemaMaxRetries() @@ -270,10 +322,11 @@ func (r *Runner) handleSchemaWithRetry( if err == nil && data != nil { elapsed := int(time.Since(startTime).Milliseconds()) - turns, sid, msgs := accumulateMetrics(allRaws) + cost, turns, sid, msgs := accumulateMetrics(allRaws) return &Result{ Result: initialRaw.Result, Parsed: dest, + CostUSD: cost, NumTurns: turns, DurationMS: elapsed, SessionID: sid, @@ -289,7 +342,7 @@ func (r *Runner) handleSchemaWithRetry( } if initialRaw.IsError && !fileExists(outputPath) && !retryableFailures[initialRaw.FailureType] { elapsed := int(time.Since(startTime).Milliseconds()) - turns, sid, msgs := accumulateMetrics(allRaws) + cost, turns, sid, msgs := accumulateMetrics(allRaws) providerError := initialRaw.ErrorMessage if providerError == "" { providerError = "Provider execution failed." @@ -299,6 +352,7 @@ func (r *Runner) handleSchemaWithRetry( IsError: true, ErrorMessage: fmt.Sprintf("%s Output file was not created at %s.", providerError, outputPath), FailureType: initialRaw.FailureType, + CostUSD: cost, NumTurns: turns, DurationMS: elapsed, SessionID: sid, @@ -317,11 +371,12 @@ func (r *Runner) handleSchemaWithRetry( case <-ctx.Done(): timer.Stop() elapsed := int(time.Since(startTime).Milliseconds()) - turns, sid, msgs := accumulateMetrics(allRaws) + cost, turns, sid, msgs := accumulateMetrics(allRaws) return &Result{ IsError: true, ErrorMessage: "context cancelled during schema retry", FailureType: FailureTimeout, + CostUSD: cost, NumTurns: turns, DurationMS: elapsed, SessionID: sid, @@ -334,6 +389,20 @@ func (r *Runner) handleSchemaWithRetry( var retryPrompt string if isCrash { retryPrompt = originalPrompt + } else if useIncremental { + // Incremental mode: patch only the fields that are missing or + // invalid, one at a time, instead of regenerating the whole + // object. The partial output file persists on disk between + // attempts, so the agent edits it in place. The original goal is + // prepended so the agent keeps the task in view (parity with the + // Python incremental retry path). + fieldErrors := DiagnoseFieldFailures(outputPath, schema, dest) + followup := BuildIncrementalFollowup(fieldErrors, outputDir, schema) + if originalPrompt != "" { + retryPrompt = originalPrompt + "\n\n" + followup + } else { + retryPrompt = followup + } } else { errorDetail := DiagnoseOutputFailure(outputPath, schema) retryPrompt = BuildFollowupPrompt(errorDetail, outputDir, schema) @@ -375,11 +444,12 @@ func (r *Runner) handleSchemaWithRetry( if err == nil && data != nil { elapsed := int(time.Since(startTime).Milliseconds()) - turns, sid, msgs := accumulateMetrics(allRaws) + cost, turns, sid, msgs := accumulateMetrics(allRaws) r.Logger.Printf("Schema validation succeeded on retry %d", retryNum+1) return &Result{ Result: retryRaw.Result, Parsed: dest, + CostUSD: cost, NumTurns: turns, DurationMS: elapsed, SessionID: sid, @@ -389,7 +459,7 @@ func (r *Runner) handleSchemaWithRetry( } elapsed := int(time.Since(startTime).Milliseconds()) - turns, sid, msgs := accumulateMetrics(allRaws) + cost, turns, sid, msgs := accumulateMetrics(allRaws) finalDiagnosis := DiagnoseOutputFailure(outputPath, schema) return &Result{ Result: allRaws[len(allRaws)-1].Result, @@ -399,6 +469,7 @@ func (r *Runner) handleSchemaWithRetry( maxRetries, finalDiagnosis, ), FailureType: FailureSchema, + CostUSD: cost, NumTurns: turns, DurationMS: elapsed, SessionID: sid, @@ -406,8 +477,20 @@ func (r *Runner) handleSchemaWithRetry( } } -func accumulateMetrics(raws []*RawResult) (totalTurns int, sessionID string, allMessages []map[string]any) { +// accumulateMetrics sums metrics across every provider execution that +// contributed to a result — including failed retry attempts, mirroring the +// Python _accumulate_metrics. CostUSD is summed only over executions that +// reported a cost; the returned pointer is nil when none did (distinguishing +// "unknown" from "$0.00"). +func accumulateMetrics(raws []*RawResult) (totalCost *float64, totalTurns int, sessionID string, allMessages []map[string]any) { for _, raw := range raws { + if raw.Metrics.CostUSD != nil { + if totalCost == nil { + zero := 0.0 + totalCost = &zero + } + *totalCost += *raw.Metrics.CostUSD + } totalTurns += raw.Metrics.NumTurns if raw.Metrics.SessionID != "" { sessionID = raw.Metrics.SessionID diff --git a/sdk/go/harness/runner_invariant_test.go b/sdk/go/harness/runner_invariant_test.go index 154008dc3..20157e06b 100644 --- a/sdk/go/harness/runner_invariant_test.go +++ b/sdk/go/harness/runner_invariant_test.go @@ -48,10 +48,10 @@ func TestInvariant_Runner_RetryCountBounded(t *testing.T) { maxRetries int wantMax int }{ - {"default max retries (3)", 0, 4}, // default 3 retries = 4 total attempts (1 + 3) - {"max retries = 1", 1, 2}, // 1 retry = 2 total attempts - {"max retries = 2", 2, 3}, // 2 retries = 3 total attempts - {"max retries = 5", 5, 6}, // 5 retries = 6 total attempts + {"default max retries (3)", 0, 4}, // default 3 retries = 4 total attempts (1 + 3) + {"max retries = 1", 1, 2}, // 1 retry = 2 total attempts + {"max retries = 2", 2, 3}, // 2 retries = 3 total attempts + {"max retries = 5", 5, 6}, // 5 retries = 6 total attempts } for _, tt := range tests { @@ -222,7 +222,7 @@ func TestInvariant_Runner_SchemaMaxRetriesBounded(t *testing.T) { schemaMaxRetries int wantMaxCalls int }{ - {"default (2)", 0, 2}, // default schemaMaxRetries = 2 + {"default (2)", 0, 2}, // default schemaMaxRetries = 2 {"1 retry", 1, 1}, {"3 retries", 3, 3}, } @@ -267,7 +267,7 @@ func TestInvariant_Runner_SchemaMaxRetriesBounded(t *testing.T) { result := runner.handleSchemaWithRetry( context.Background(), initialRaw, schema, &dest, dir, - time.Now(), prov, opts, "test prompt", + time.Now(), prov, opts, "test prompt", false, ) assert.True(t, result.IsError, "should fail after exhausting retries") diff --git a/sdk/go/harness/runner_test.go b/sdk/go/harness/runner_test.go index 804346ed2..e17469039 100644 --- a/sdk/go/harness/runner_test.go +++ b/sdk/go/harness/runner_test.go @@ -101,7 +101,7 @@ func TestRunner_HandleSchemaWithRetry_FirstAttemptSuccess(t *testing.T) { raw := &RawResult{Result: "done", Metrics: Metrics{NumTurns: 1}} result := runner.handleSchemaWithRetry( context.Background(), raw, schema, &dest, dir, - time.Now(), mock, Options{Provider: "opencode"}, "test prompt", + time.Now(), mock, Options{Provider: "opencode"}, "test prompt", false, ) assert.False(t, result.IsError) @@ -136,7 +136,7 @@ func TestRunner_HandleSchemaWithRetry_StdoutFallback(t *testing.T) { result := runner.handleSchemaWithRetry( context.Background(), raw, schema, &dest, dir, - time.Now(), mock, Options{Provider: "opencode"}, "test prompt", + time.Now(), mock, Options{Provider: "opencode"}, "test prompt", false, ) assert.False(t, result.IsError) @@ -196,7 +196,7 @@ func TestRunner_HandleSchemaWithRetry_RetrySuccess(t *testing.T) { result := runner.handleSchemaWithRetry( context.Background(), initialRaw, schema, &dest, dir, - time.Now(), mock2, Options{Provider: "opencode", SchemaMaxRetries: 2}, "test prompt", + time.Now(), mock2, Options{Provider: "opencode", SchemaMaxRetries: 2}, "test prompt", false, ) assert.False(t, result.IsError) @@ -250,7 +250,7 @@ func TestRunner_HandleSchemaWithRetry_AllRetriesFail(t *testing.T) { result := runner.handleSchemaWithRetry( context.Background(), initialRaw, schema, &dest, dir, - time.Now(), mock, Options{Provider: "opencode", SchemaMaxRetries: 2}, "test prompt", + time.Now(), mock, Options{Provider: "opencode", SchemaMaxRetries: 2}, "test prompt", false, ) assert.True(t, result.IsError) diff --git a/sdk/go/harness/schema.go b/sdk/go/harness/schema.go index 7be41dba0..2c0340c56 100644 --- a/sdk/go/harness/schema.go +++ b/sdk/go/harness/schema.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "reflect" "regexp" "sort" "strings" @@ -352,3 +353,224 @@ func BuildFollowupPrompt(errorMessage string, dir string, jsonSchema map[string] return b.String() } + +// topLevelField describes one top-level schema property for the incremental +// build: its name and whether the schema marks it required. +type topLevelField struct { + Name string + Required bool +} + +// getTopLevelFields returns the schema's top-level properties and whether each +// is required. Mirrors the Python _schema.get_top_level_fields. +// +// Note on ordering: Go's JSON schema is a map[string]any, so the original +// property order from the source schema is not preserved (Python dicts keep +// insertion order). Field names are sorted for deterministic prompt output; +// this affects only the order of the field list, not correctness — the agent +// is asked to add every field regardless of order. +func getTopLevelFields(jsonSchema map[string]any) []topLevelField { + props, _ := jsonSchema["properties"].(map[string]any) + required := requiredSet(jsonSchema) + + names := make([]string, 0, len(props)) + for name := range props { + names = append(names, name) + } + sort.Strings(names) + + fields := make([]topLevelField, 0, len(names)) + for _, name := range names { + fields = append(fields, topLevelField{Name: name, Required: required[name]}) + } + return fields +} + +// requiredSet extracts the schema's "required" list as a set. +func requiredSet(jsonSchema map[string]any) map[string]bool { + set := make(map[string]bool) + if req, ok := jsonSchema["required"].([]any); ok { + for _, r := range req { + if s, ok := r.(string); ok { + set[s] = true + } + } + } + // Also accept a pre-typed []string (callers who build schemas in Go). + if req, ok := jsonSchema["required"].([]string); ok { + for _, s := range req { + set[s] = true + } + } + return set +} + +// BuildIncrementalPromptSuffix builds the OUTPUT REQUIREMENTS suffix that +// instructs a field-by-field build. Byte-for-byte parity with the Python +// _schema.build_incremental_prompt_suffix. +func BuildIncrementalPromptSuffix(jsonSchema map[string]any, dir string) string { + outputPath := OutputPath(dir) + schemaJSON, err := json.MarshalIndent(jsonSchema, "", " ") + if err != nil { + // Fall back to the single-shot suffix if the schema cannot be + // serialized — matches the defensive posture of BuildPromptSuffix. + return BuildPromptSuffix(jsonSchema, dir) + } + + fields := getTopLevelFields(jsonSchema) + fieldLineParts := make([]string, 0, len(fields)) + for _, f := range fields { + req := "optional" + if f.Required { + req = "required" + } + fieldLineParts = append(fieldLineParts, fmt.Sprintf(" - %s (%s)", f.Name, req)) + } + fieldLines := strings.Join(fieldLineParts, "\n") + + var schemaRef string + if estimateTokens(string(schemaJSON)) > largeSchemaTokenThreshold { + schemaPath := SchemaPath(dir) + _ = writeSchemaFile(string(schemaJSON), dir) + schemaRef = fmt.Sprintf( + "The full JSON Schema is at: %s\nRead it for each field's exact shape.\n", + schemaPath, + ) + } else { + schemaRef = fmt.Sprintf("Full JSON Schema:\n%s\n", string(schemaJSON)) + } + + return "\n\n---\n" + + "CRITICAL OUTPUT REQUIREMENTS (incremental build):\n" + + fmt.Sprintf("Produce a single JSON object in this file using your Write/Edit tools: %s\n", outputPath) + + "Build it ONE FIELD AT A TIME so nothing gets truncated:\n" + + " 1. First create the file with an empty object: {}\n" + + " 2. Then add each field listed below one at a time using Edit, and after\n" + + " each edit re-read the file to confirm it is still valid JSON.\n" + + " 3. Each field's value MUST conform to its shape in the schema.\n" + + " 4. Do not finish until every required field is present.\n\n" + + fmt.Sprintf("Top-level fields to add:\n%s\n\n", fieldLines) + + fmt.Sprintf("%s\n", schemaRef) + + fmt.Sprintf("The final file at %s MUST contain ONLY the complete valid JSON "+ + "object — no markdown fences, no commentary, no extra text.", outputPath) +} + +// DiagnoseFieldFailures maps each missing/invalid top-level field to a short +// reason. Returns an empty map when the file validates cleanly. Mirrors the +// Python _schema.diagnose_field_failures, adapted to the Go dest-struct +// contract. +// +// Go has no per-field pydantic error list, so where Python enumerates +// validation errors by field location, Go can only detect (a) missing required +// fields — via the schema's "required" list against the parsed object — and +// (b) a whole-document type mismatch when the object fails to unmarshal into +// dest, reported once under the "_root" key (matching Python's fallback). +func DiagnoseFieldFailures(filePath string, jsonSchema map[string]any, dest any) map[string]string { + props, _ := jsonSchema["properties"].(map[string]any) + propNames := make([]string, 0, len(props)) + for name := range props { + propNames = append(propNames, name) + } + sort.Strings(propNames) + required := requiredSet(jsonSchema) + + data, err := ReadAndParse(filePath) + if err != nil { + data, err = ReadRepairAndParse(filePath) + } + + failures := make(map[string]string) + + if err != nil || data == nil { + // Whole file unusable — report every required field (or all fields if + // none are required) as needing to be written. + targets := make([]string, 0) + for _, name := range propNames { + if required[name] { + targets = append(targets, name) + } + } + if len(targets) == 0 { + targets = propNames + } + for _, name := range targets { + failures[name] = "output file missing or not a JSON object" + } + return failures + } + + // Required-field presence check (sorted for deterministic prompt output). + reqNames := make([]string, 0, len(required)) + for name := range required { + reqNames = append(reqNames, name) + } + sort.Strings(reqNames) + for _, name := range reqNames { + if _, present := data[name]; !present { + failures[name] = "missing required field" + } + } + + // Validation against the destination struct. Use a throwaway instance so + // the caller's dest is not mutated during diagnosis. + if dest != nil { + if fresh := freshDest(dest); fresh != nil { + if e := unmarshalInto(data, fresh); e != nil { + if _, exists := failures["_root"]; !exists { + failures["_root"] = truncate(e.Error(), 200) + } + } + } + } + + return failures +} + +// freshDest returns a new zero-valued instance of the same type dest points to +// (a pointer), so validation can run without mutating the caller's value. +// Returns nil when dest is not a non-nil pointer. +func freshDest(dest any) any { + rv := reflect.ValueOf(dest) + if rv.Kind() != reflect.Pointer || rv.IsNil() { + return nil + } + return reflect.New(rv.Elem().Type()).Interface() +} + +// BuildIncrementalFollowup builds the follow-up prompt that asks the agent to +// patch only the failing fields. Byte-for-byte parity with the Python +// _schema.build_incremental_followup. +func BuildIncrementalFollowup(fieldErrors map[string]string, dir string, jsonSchema map[string]any) string { + outputPath := OutputPath(dir) + + // Deterministic order (Go maps randomize; Python preserves dict order). + names := make([]string, 0, len(fieldErrors)) + for name := range fieldErrors { + names = append(names, name) + } + sort.Strings(names) + lineParts := make([]string, 0, len(names)) + for _, name := range names { + lineParts = append(lineParts, fmt.Sprintf(" - %s: %s", name, fieldErrors[name])) + } + fieldLines := strings.Join(lineParts, "\n") + + var schemaRef string + schemaJSON, err := json.MarshalIndent(jsonSchema, "", " ") + if err == nil && estimateTokens(string(schemaJSON)) > largeSchemaTokenThreshold { + schemaPath := SchemaPath(dir) + if _, statErr := os.Stat(schemaPath); statErr != nil { + _ = writeSchemaFile(string(schemaJSON), dir) + } + schemaRef = fmt.Sprintf("Full schema is at: %s\n", schemaPath) + } else if err == nil { + schemaRef = fmt.Sprintf("Full schema:\n%s\n", string(schemaJSON)) + } + + return fmt.Sprintf( + "PARTIAL OUTPUT NEEDS FIXES. The JSON at %s is incomplete or invalid.\n", outputPath) + + "Patch ONLY these fields, one at a time, using Edit, keeping the file valid JSON after each change:\n" + + fmt.Sprintf("%s\n\n", fieldLines) + + schemaRef + + "Leave every already-correct field unchanged. Do NOT rewrite the whole file." +}