iop/apps/edge/internal/service/single_request_tool_loop_test.go
toki dc9a9a8c59 feat(agent): 단일 요청 Agent 실행 경계를 구현한다
승인된 execution preset을 Edge 조정 경계와 Node workspace/tool 실행 경계로 연결해 단일 요청 수명주기와 관측 계약을 일관되게 처리한다.
2026-08-07 07:03:55 +09:00

443 lines
19 KiB
Go

package service
import (
"context"
"encoding/json"
"errors"
"strings"
"sync/atomic"
"testing"
"time"
toki "git.toki-labs.com/toki/proto-socket/go"
edgenode "iop/apps/edge/internal/node"
"iop/packages/go/config"
iop "iop/proto/gen/iop"
)
type scriptedInternalToolExecutor struct {
calls []InternalWorkspaceToolCall
results chan InternalWorkspaceToolResult
seenResults []InternalWorkspaceToolResult
continueCount atomic.Int32
}
func newScriptedInternalToolExecutor(calls ...InternalWorkspaceToolCall) *scriptedInternalToolExecutor {
return &scriptedInternalToolExecutor{calls: calls, results: make(chan InternalWorkspaceToolResult, len(calls)+1)}
}
func (e *scriptedInternalToolExecutor) ExecuteSingleRequest(ctx context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
sequence := uint64(1)
if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, sequence, SingleRequestStatePlanning)); err != nil {
return err
}
for index := range e.calls {
call := e.calls[index].Clone()
if call.RequestID == "" {
call.RequestID = req.RequestID
}
if call.StageID == "" {
call.StageID = "plan"
}
sequence++
if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{
RequestID: req.RequestID, Sequence: sequence,
Stage: SingleRequestStateInternalTool, SavedStage: SingleRequestStatePlanning,
ToolCall: call,
}); err != nil {
return err
}
select {
case result := <-e.results:
e.seenResults = append(e.seenResults, result.Clone())
case <-ctx.Done():
return ctx.Err()
}
sequence++
if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{
RequestID: req.RequestID, Sequence: sequence,
Stage: SingleRequestStatePlanning, SavedStage: SingleRequestStatePlanning,
}); err != nil {
return err
}
}
for _, stage := range []SingleRequestState{SingleRequestStateWorking, SingleRequestStateReviewing, SingleRequestStateFinalizing} {
sequence++
envelope := testEnvelope(req.RequestID, sequence, stage)
if stage == SingleRequestStateFinalizing {
envelope.Result = &SingleRequestResult{Output: "private tools completed"}
}
if err := ctrl.SubmitEnvelope(envelope); err != nil {
return err
}
}
return nil
}
func (e *scriptedInternalToolExecutor) ContinueInternalTool(_ context.Context, result InternalWorkspaceToolResult) error {
e.continueCount.Add(1)
e.results <- result.Clone()
return nil
}
func internalLoopWorkspace() config.WorkspaceDefinition {
return config.WorkspaceDefinition{
Ref: "workspace-loop", Platform: "darwin", Root: "/Users/operator/project",
Operations: []config.WorkspaceOperation{
config.WorkspaceOpRead, config.WorkspaceOpList, config.WorkspaceOpWrite,
config.WorkspaceOpDelete, config.WorkspaceOpCommand,
},
Commands: []config.WorkspaceCommandDefinition{{ID: "test"}},
EnvironmentAllowlist: []string{"IOP_MODE"},
MaxReadBytes: 1024, MaxWriteBytes: 1024, MaxOutputBytes: 1024, MaxCommandTimeoutMS: 1000,
}
}
func newInternalToolLoopService(t *testing.T, executor SingleRequestExecutor) (*Service, *toki.TcpClient) {
t.Helper()
edgeClient, nodeClient := workspaceWirePipe(t)
registry := edgenode.NewRegistry()
registry.Register(&edgenode.NodeEntry{NodeID: "node-loop", Client: edgeClient})
service := New(registry, nil)
service.SetNodeStore(workspaceStore("node-loop", internalLoopWorkspace()))
service.SetSingleRequestExecutor(executor)
return service, nodeClient
}
func TestSingleRequestInternalToolLoopRequiresOptionalContinuation(t *testing.T) {
executor := &channelFakeExecutor{fn: func(_ context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, 1, SingleRequestStatePlanning)); err != nil {
return err
}
return ctrl.SubmitEnvelope(SingleRequestEnvelope{
RequestID: req.RequestID, Sequence: 2,
Stage: SingleRequestStateInternalTool, SavedStage: SingleRequestStatePlanning,
ToolCall: &InternalWorkspaceToolCall{
RequestID: req.RequestID, StageID: "plan", ToolCallID: "tool-1",
Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`),
},
})
}}
service, node := newInternalToolLoopService(t, executor)
var openCount atomic.Int32
installInternalLoopOpenResponder(node, &openCount)
handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
if err != nil {
t.Fatalf("StartSingleRequest: %v", err)
}
if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestInternalToolUnavailable) {
t.Fatalf("Wait error = %v, want unavailable continuation", err)
}
if openCount.Load() != 0 {
t.Fatalf("unavailable continuation opened workspace %d times", openCount.Load())
}
}
func internalLoopRequest(t *testing.T, mutate func(*SingleRequestBinding)) SingleRequestRequest {
t.Helper()
binding := createTestBinding(t)
binding.WorkspaceRef = "workspace-loop"
if mutate != nil {
mutate(binding)
}
return SingleRequestRequest{RequestID: "request-loop", Binding: binding, Prompt: "complete the private task"}
}
func installInternalLoopOpenResponder(node *toki.TcpClient, count *atomic.Int32) {
toki.AddRequestListenerTyped[*iop.WorkspaceOpenRequest, *iop.WorkspaceOpenResponse](&node.Communicator, func(req *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error) {
count.Add(1)
return &iop.WorkspaceOpenResponse{
RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(),
Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
}, nil
})
toki.AddRequestListenerTyped[*iop.WorkspaceCleanupRequest, *iop.WorkspaceCleanupResponse](&node.Communicator, func(req *iop.WorkspaceCleanupRequest) (*iop.WorkspaceCleanupResponse, error) {
return &iop.WorkspaceCleanupResponse{RequestId: req.GetRequestId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
})
}
func TestSingleRequestInternalToolLoopMultipleOrdered(t *testing.T) {
executor := newScriptedInternalToolExecutor(
InternalWorkspaceToolCall{ToolCallID: "tool-read", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
InternalWorkspaceToolCall{ToolCallID: "tool-write", Name: InternalWorkspaceToolWrite, Arguments: json.RawMessage(`{"relative_path":"result.txt","content":"done"}`)},
InternalWorkspaceToolCall{ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"test","environment":{"IOP_MODE":"safe"}}`)},
)
service, node := newInternalToolLoopService(t, executor)
var openCount atomic.Int32
installInternalLoopOpenResponder(node, &openCount)
toolOrder := make(chan string, 3)
commandTimeout := make(chan int64, 1)
toki.AddRequestListenerTyped[*iop.WorkspaceToolRequest, *iop.WorkspaceToolResponse](&node.Communicator, func(req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
toolOrder <- req.GetToolCallId()
response := &iop.WorkspaceToolResponse{
RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(),
Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
}
switch req.GetOperation() {
case iop.WorkspaceOperation_WORKSPACE_OPERATION_READ:
response.Content = []byte("source")
case iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND:
commandTimeout <- req.GetTimeoutMs()
response.Stdout = []byte("ok")
}
return response, nil
})
handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
if err != nil {
t.Fatalf("StartSingleRequest: %v", err)
}
waitForState(t, handle, SingleRequestStateFinalizing)
waitForSingleRequestCleanup(t, handle)
if err := handle.AcknowledgeTerminal(true); err != nil {
t.Fatalf("AcknowledgeTerminal: %v", err)
}
result, err := waitForExecution(t, handle)
if err != nil || result.Output != "private tools completed" {
t.Fatalf("Wait = (%q, %v)", result.Output, err)
}
if openCount.Load() != 1 || executor.continueCount.Load() != 3 || len(executor.seenResults) != 3 {
t.Fatalf("open=%d continuations=%d results=%d", openCount.Load(), executor.continueCount.Load(), len(executor.seenResults))
}
for index, want := range []string{"tool-read", "tool-write", "tool-command"} {
if got := <-toolOrder; got != want {
t.Fatalf("tool order[%d]=%q, want %q", index, got, want)
}
if executor.seenResults[index].ToolCallID != want || executor.seenResults[index].StageID != "plan" {
t.Fatalf("correlated result[%d]=%+v", index, executor.seenResults[index])
}
}
if got := <-commandTimeout; got <= 0 || got > int64(internalLoopWorkspace().MaxCommandTimeoutMS) {
t.Fatalf("command timeout = %d, want within workspace maximum %d", got, internalLoopWorkspace().MaxCommandTimeoutMS)
}
for progress := range handle.Progress() {
if progress.Message == InternalWorkspaceToolRead || progress.Message == InternalWorkspaceToolWrite || progress.Message == InternalWorkspaceToolCommand {
t.Fatalf("internal tool protocol reached progress: %+v", progress)
}
}
}
func waitForSingleRequestCleanup(t *testing.T, handle SingleRequestExecution) {
t.Helper()
internal, ok := handle.(*singleRequestHandle)
if !ok {
t.Fatal("execution does not expose coordinator cleanup state")
}
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
internal.mu.Lock()
complete := internal.cleanupComplete
internal.mu.Unlock()
if complete {
return
}
time.Sleep(time.Millisecond)
}
t.Fatal("workspace cleanup did not complete")
}
func TestSingleRequestInternalToolLoopFailsClosed(t *testing.T) {
const rawSentinel = "RAW-TOOL-SENTINEL"
tests := []struct {
name string
calls []InternalWorkspaceToolCall
mutateBinding func(*SingleRequestBinding)
respond func(*iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse
want error
wantWireCalls int32
wantContinuations int32
}{
{
name: "identity mismatch",
calls: []InternalWorkspaceToolCall{{RequestID: "other-request", StageID: "plan", ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
want: ErrSingleRequestIdentityMismatch,
},
{
name: "malformed arguments",
calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md","raw":"` + rawSentinel + `"}`)}},
want: ErrSingleRequestInternalToolInvalidCall,
},
{
name: "capability denied",
calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"missing"}`)}},
want: ErrSingleRequestInternalToolDenied,
},
{
name: "environment capability denied",
calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"test","environment":{"NOT_ALLOWED":"value"}}`)}},
want: ErrSingleRequestInternalToolDenied,
},
{
name: "duplicate tool id",
calls: []InternalWorkspaceToolCall{
{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
},
want: ErrSingleRequestInternalToolInvalidCall, wantWireCalls: 1, wantContinuations: 1,
},
{
name: "iteration budget",
calls: []InternalWorkspaceToolCall{
{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
{ToolCallID: "tool-2", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
},
mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxToolIterations = 1 },
want: ErrSingleRequestInternalToolBudget, wantWireCalls: 1, wantContinuations: 1,
},
{
name: "stale Node response",
calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
respond: func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: "stale-tool", Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
},
want: ErrSingleRequestInternalToolFailed, wantWireCalls: 1,
},
{
name: "output budget",
calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxOutputBytes = 3 },
respond: func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, Content: []byte("four")}
},
want: ErrSingleRequestInternalToolBudget, wantWireCalls: 1,
},
{
name: "cumulative output budget",
calls: []InternalWorkspaceToolCall{
{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
{ToolCallID: "tool-2", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
},
mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxOutputBytes = 3 },
respond: func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, Content: []byte("xx")}
},
want: ErrSingleRequestInternalToolBudget, wantWireCalls: 2, wantContinuations: 1,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
executor := newScriptedInternalToolExecutor(test.calls...)
service, node := newInternalToolLoopService(t, executor)
var openCount atomic.Int32
installInternalLoopOpenResponder(node, &openCount)
var toolCount atomic.Int32
toki.AddRequestListenerTyped[*iop.WorkspaceToolRequest, *iop.WorkspaceToolResponse](&node.Communicator, func(req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
toolCount.Add(1)
if test.respond != nil {
return test.respond(req), nil
}
return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
})
handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, test.mutateBinding))
if err != nil {
t.Fatalf("StartSingleRequest: %v", err)
}
_, err = waitForExecution(t, handle)
if !errors.Is(err, test.want) {
t.Fatalf("Wait error = %v, want %v", err, test.want)
}
if strings.Contains(err.Error(), rawSentinel) {
t.Fatalf("raw call leaked in error %q", err)
}
wantOpenCalls := int32(0)
if test.wantWireCalls > 0 {
wantOpenCalls = 1
}
if openCount.Load() != wantOpenCalls || toolCount.Load() != test.wantWireCalls || executor.continueCount.Load() != test.wantContinuations {
t.Fatalf("open=%d wire calls=%d continuations=%d, want %d/%d/%d", openCount.Load(), toolCount.Load(), executor.continueCount.Load(), wantOpenCalls, test.wantWireCalls, test.wantContinuations)
}
})
}
}
func TestSingleRequestInternalToolLoopCancelPropagates(t *testing.T) {
executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand,
Arguments: json.RawMessage(`{"command_id":"test"}`),
})
service, node := newInternalToolLoopService(t, executor)
var openCount atomic.Int32
installInternalLoopOpenResponder(node, &openCount)
var sequence atomic.Int32
toolEntered := make(chan struct{})
release := make(chan struct{})
cancelReached := make(chan *iop.WorkspaceCancelRequest, 1)
serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
close(toolEntered)
<-release
return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
})
serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceCancelRequest) *iop.WorkspaceCancelResponse {
cancelReached <- req
return &iop.WorkspaceCancelResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
})
handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
if err != nil {
t.Fatalf("StartSingleRequest: %v", err)
}
select {
case <-toolEntered:
case <-time.After(2 * time.Second):
close(release)
t.Fatal("tool did not reach Node")
}
handle.Cancel()
select {
case cancel := <-cancelReached:
if cancel.GetRequestId() != "request-loop" || cancel.GetStageId() != "plan" || cancel.GetToolCallId() != "tool-command" {
close(release)
t.Fatalf("cancel identity = %+v", cancel)
}
case <-time.After(2 * time.Second):
close(release)
t.Fatal("typed cancel did not reach Node")
}
close(release)
if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestCancelled) {
t.Fatalf("Wait error = %v, want cancelled", err)
}
if executor.continueCount.Load() != 0 {
t.Fatalf("cancelled call delivered %d continuations", executor.continueCount.Load())
}
}
func TestSingleRequestInternalToolLoopStageDeadline(t *testing.T) {
executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand,
Arguments: json.RawMessage(`{"command_id":"test"}`),
})
service, node := newInternalToolLoopService(t, executor)
var openCount atomic.Int32
installInternalLoopOpenResponder(node, &openCount)
var sequence atomic.Int32
toolEntered := make(chan struct{})
release := make(chan struct{})
serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
close(toolEntered)
<-release
return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
})
handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, func(binding *SingleRequestBinding) {
binding.Limits.StageTimeoutMS = 50
}))
if err != nil {
t.Fatalf("StartSingleRequest: %v", err)
}
select {
case <-toolEntered:
case <-time.After(2 * time.Second):
close(release)
t.Fatal("tool did not reach Node")
}
if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestInternalToolBudget) {
close(release)
t.Fatalf("Wait error = %v, want budget exhaustion", err)
}
close(release)
if handle.State() != SingleRequestStateFailed || executor.continueCount.Load() != 0 {
t.Fatalf("state=%s continuations=%d, want failed/0", handle.State(), executor.continueCount.Load())
}
}