diff --git a/pkg/agent/approval.go b/pkg/agent/approval.go index b3d7b32ef..4dc573492 100644 --- a/pkg/agent/approval.go +++ b/pkg/agent/approval.go @@ -29,8 +29,8 @@ type ( } ApprovalResult struct { - Approved bool - Message string + Approved bool `json:"approved"` + Message string `json:"message"` } ResumeInput struct { diff --git a/pkg/agent/checkpoint.go b/pkg/agent/checkpoint.go new file mode 100644 index 000000000..952628e88 --- /dev/null +++ b/pkg/agent/checkpoint.go @@ -0,0 +1,80 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package agent + +import ( + "context" + + "go.probo.inc/probo/pkg/llm" +) + +type ( + CheckpointStatus string + + Checkpoint struct { + Version int `json:"version"` + Status CheckpointStatus `json:"status"` + AgentName string `json:"agent_name"` + Messages []llm.Message `json:"messages"` + Usage llm.Usage `json:"usage"` + Turns int `json:"turns"` + ToolUsedInRun bool `json:"tool_used_in_run"` + + // Approval-interrupted checkpoints carry pending tool calls. + PendingToolCalls []llm.ToolCall `json:"pending_tool_calls,omitempty"` + PendingApprovals []llm.ToolCall `json:"pending_approvals,omitempty"` + ApprovalInput map[string]ApprovalResult `json:"approval_input,omitempty"` + + // Nested agent-as-tool suspension: one entry per suspended inner agent. + AllToolCalls []llm.ToolCall `json:"all_tool_calls,omitempty"` + InnerCheckpoints map[string]*Checkpoint `json:"inner_checkpoints,omitempty"` + CompletedCalls []CompletedCall `json:"completed_calls,omitempty"` + } + + CompletedCall struct { + ToolCallID string `json:"tool_call_id"` + Result ToolResult `json:"result"` + } + + // CheckpointStore is supervisor-internal. Implementations may use raw + // run IDs because public API/service methods perform tenant scoping and + // authorization before a run reaches the supervisor. + CheckpointStore interface { + Save(ctx context.Context, runID string, cp *Checkpoint) error + Load(ctx context.Context, runID string) (*Checkpoint, error) + Delete(ctx context.Context, runID string) error + } + + AgentRegistry interface { + Agent(name string) (*Agent, error) + } + + SuspendedError struct { + RunID string // Set when the outer loop has a store+runID (supervisor-managed). + Checkpoint *Checkpoint // Set when returning from an inner agent-as-tool (no store). + } +) + +const ( + CheckpointVersion = 1 + MaxCheckpointBytes = 10 * 1024 * 1024 + + CheckpointStatusSuspended CheckpointStatus = "suspended" + CheckpointStatusAwaitingApproval CheckpointStatus = "awaiting_approval" +) + +func (e *SuspendedError) Error() string { + return "agent run suspended" +} diff --git a/pkg/agent/errors.go b/pkg/agent/errors.go index 3f44fec58..06ed0148a 100644 --- a/pkg/agent/errors.go +++ b/pkg/agent/errors.go @@ -59,7 +59,7 @@ type ( inner *InterruptedError toolCallID string allToolCalls []llm.ToolCall - completedCalls []completedCall + completedCalls []CompletedCall } outerLoopState struct { @@ -69,14 +69,9 @@ type ( turns int allToolCalls []llm.ToolCall toolCallID string - completedCalls []completedCall + completedCalls []CompletedCall innerInterrupt *InterruptedError } - - completedCall struct { - toolCallID string - result ToolResult - } ) func (e *MaxTurnsExceededError) Error() string { diff --git a/pkg/agent/stop.go b/pkg/agent/stop.go new file mode 100644 index 000000000..558a8bf64 --- /dev/null +++ b/pkg/agent/stop.go @@ -0,0 +1,31 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package agent + +import "context" + +type stopSignalKey struct{} + +// WithStopSignal returns a derived context carrying a stop channel. +// The agent loop checks this channel at each turn boundary. +func WithStopSignal(ctx context.Context, ch <-chan struct{}) context.Context { + return context.WithValue(ctx, stopSignalKey{}, ch) +} + +// stopSignalFrom retrieves the stop channel from the context, or nil. +func stopSignalFrom(ctx context.Context) <-chan struct{} { + ch, _ := ctx.Value(stopSignalKey{}).(<-chan struct{}) + return ch +} diff --git a/pkg/agent/stop_test.go b/pkg/agent/stop_test.go new file mode 100644 index 000000000..49f8d860f --- /dev/null +++ b/pkg/agent/stop_test.go @@ -0,0 +1,63 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package agent + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestStopSignal(t *testing.T) { + t.Parallel() + + t.Run("absent returns nil", func(t *testing.T) { + t.Parallel() + assert.Nil(t, stopSignalFrom(context.Background())) + }) + + t.Run("round trip", func(t *testing.T) { + t.Parallel() + ch := make(chan struct{}) + ctx := WithStopSignal(context.Background(), ch) + assert.Equal(t, (<-chan struct{})(ch), stopSignalFrom(ctx)) + }) + + t.Run("non-blocking when open", func(t *testing.T) { + t.Parallel() + ch := make(chan struct{}) + ctx := WithStopSignal(context.Background(), ch) + sig := stopSignalFrom(ctx) + select { + case <-sig: + t.Fatal("should not fire") + default: + } + }) + + t.Run("fires when closed", func(t *testing.T) { + t.Parallel() + ch := make(chan struct{}) + ctx := WithStopSignal(context.Background(), ch) + close(ch) + sig := stopSignalFrom(ctx) + select { + case <-sig: + default: + t.Fatal("should have fired") + } + }) +} diff --git a/pkg/agent/tool.go b/pkg/agent/tool.go index c898aaf0e..471eb0654 100644 --- a/pkg/agent/tool.go +++ b/pkg/agent/tool.go @@ -25,8 +25,8 @@ import ( type ( ToolResult struct { - Content string - IsError bool + Content string `json:"content"` + IsError bool `json:"is_error"` } // ToolDescriptor describes a tool's name and LLM definition.