Add files via upload

2026-08-05 02:18:46 +02:00 · 2026-07-21 20:52:43 +08:00
parent c925c2e74b
commit bfe1d28650
10 changed files with 1916 additions and 227 deletions
@@ -3,7 +3,9 @@ package mcp
 import (
 	"context"
 	"errors"
+	"strings"
 	"testing"
+	"time"

 	"cyberstrike-ai/internal/authctx"

@@ -34,3 +36,147 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
 		t.Fatalf("execution owner = %#v, want u1", execution)
 	}
 }
+
+func TestServerCallToolBoundedWaitForInternalTool(t *testing.T) {
+	server := NewServer(zap.NewNop())
+	server.toolWaitTimeout = 10 * time.Millisecond
+	release := make(chan struct{})
+	started := make(chan struct{})
+	server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
+		close(started)
+		select {
+		case <-release:
+			return &ToolResult{Content: []Content{{Type: "text", Text: "internal done"}}}, nil
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	})
+
+	callCtx, callCancel := context.WithCancel(context.Background())
+	result, executionID, err := server.CallTool(callCtx, "slow", nil)
+	if err != nil {
+		t.Fatalf("CallTool returned error: %v", err)
+	}
+	if executionID == "" || result == nil || !result.IsError {
+		t.Fatalf("expected soft timeout with execution id, result=%#v id=%q", result, executionID)
+	}
+	if text := ToolResultPlainText(result); !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
+		t.Fatalf("timeout result missing execution guidance: %q", text)
+	}
+	select {
+	case <-started:
+	default:
+		t.Fatal("internal worker did not start")
+	}
+	callCancel()
+	close(release)
+
+	snapshot, err := server.executionService.Wait(context.Background(), executionID, time.Second)
+	if err != nil {
+		t.Fatalf("wait internal execution: %v", err)
+	}
+	if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusCompleted {
+		t.Fatalf("snapshot = %#v, want completed", snapshot)
+	}
+	if got := ToolResultPlainText(snapshot.Execution.Result); got != "internal done" {
+		t.Fatalf("result = %q, want internal done", got)
+	}
+}
+
+func TestWaitToolExecutionWaitsForInternalActiveExecution(t *testing.T) {
+	server := NewServer(zap.NewNop())
+	server.toolWaitTimeout = 10 * time.Millisecond
+	release := make(chan struct{})
+	server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
+		select {
+		case <-release:
+			return &ToolResult{Content: []Content{{Type: "text", Text: "wait saw completion"}}}, nil
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	})
+	RegisterExecutionControlTools(server, nil)
+
+	result, executionID, err := server.CallTool(context.Background(), "slow", nil)
+	if err != nil {
+		t.Fatalf("CallTool returned error: %v", err)
+	}
+	if result == nil || !result.IsError || executionID == "" {
+		t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
+	}
+
+	done := make(chan *ToolResult, 1)
+	errCh := make(chan error, 1)
+	go func() {
+		waitResult, _, waitErr := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
+			"execution_id":    executionID,
+			"timeout_seconds": 1,
+		})
+		if waitErr != nil {
+			errCh <- waitErr
+			return
+		}
+		done <- waitResult
+	}()
+
+	select {
+	case <-done:
+		t.Fatal("wait_tool_execution returned before target execution completed")
+	case err := <-errCh:
+		t.Fatalf("wait_tool_execution errored before release: %v", err)
+	case <-time.After(50 * time.Millisecond):
+	}
+	close(release)
+
+	select {
+	case err := <-errCh:
+		t.Fatalf("wait_tool_execution returned error: %v", err)
+	case waitResult := <-done:
+		if waitResult == nil || waitResult.IsError {
+			t.Fatalf("expected successful wait result, got %#v", waitResult)
+		}
+		if body := ToolResultPlainText(waitResult); !strings.Contains(body, "wait saw completion") || !strings.Contains(body, `"status": "completed"`) {
+			t.Fatalf("wait result missing completed target: %s", body)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("wait_tool_execution did not return after target completion")
+	}
+}
+
+func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) {
+	server := NewServer(zap.NewNop())
+	server.toolWaitTimeout = 10 * time.Millisecond
+	release := make(chan struct{})
+	server.RegisterTool(Tool{Name: "slow_observed", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
+		<-release
+		return &ToolResult{Content: []Content{{Type: "text", Text: "done"}}}, nil
+	})
+	RegisterExecutionControlTools(server, nil)
+
+	result, executionID, err := server.CallTool(context.Background(), "slow_observed", nil)
+	if err != nil {
+		t.Fatalf("CallTool returned error: %v", err)
+	}
+	if result == nil || !result.IsError || executionID == "" {
+		t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
+	}
+
+	waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
+		"execution_id":    executionID,
+		"timeout_seconds": 0.01,
+	})
+	if err != nil {
+		t.Fatalf("wait_tool_execution returned error: %v", err)
+	}
+	if waitResult == nil {
+		t.Fatal("missing wait result")
+	}
+	if waitResult.IsError {
+		t.Fatalf("wait timeout should be a successful observation, got %#v", waitResult)
+	}
+	body := ToolResultPlainText(waitResult)
+	if !strings.Contains(body, `"status": "running"`) || !strings.Contains(body, "本次等待已到达") {
+		t.Fatalf("wait timeout body missing running status/guidance: %s", body)
+	}
+	close(release)
+}