fix: harden debug chat auth and cancellation

onutc · onutc · commit c6cc66ab49b5 · 2026-03-24T20:24:37.000+01:00
diff --git a/api/acp_bootstrap.go b/api/acp_bootstrap.go
@@ -198,6 +198,23 @@ func (c *acpBootstrapInstanceClient) writeJSON(ctx context.Context, payload any)
 	return c.conn.WriteJSON(payload)
 }
 
+func (c *acpBootstrapInstanceClient) notify(ctx context.Context, method string, params any) error {
+	return c.writeJSON(ctx, map[string]any{
+		"jsonrpc": "2.0",
+		"method":  method,
+		"params":  params,
+	})
+}
+
+func (c *acpBootstrapInstanceClient) cancelPrompt(ctx context.Context, sessionID string) error {
+	if strings.TrimSpace(sessionID) == "" {
+		return nil
+	}
+	return c.notify(ctx, "session/cancel", map[string]any{
+		"sessionId": sessionID,
+	})
+}
+
 func (c *acpBootstrapInstanceClient) readMessage(ctx context.Context) (*acpBootstrapJSONRPCMessage, error) {
 	if deadline, ok := ctx.Deadline(); ok {
 		if err := c.conn.SetReadDeadline(deadline); err != nil {
diff --git a/api/internal_debug_chat.go b/api/internal_debug_chat.go
@@ -8,6 +8,7 @@ import (
 	"log"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -209,6 +210,16 @@ func (s *server) cleanupInternalDebugConversation(ctx context.Context, conversat
 func (s *server) runInternalDebugChatPrompt(ctx context.Context, spritz *spritzv1.Spritz, sessionID, message string) (*acpPromptResult, error) {
 	runCtx, cancel := context.WithTimeout(ctx, s.acp.promptTimeout)
 	defer cancel()
+	cancelWatcherDone := make(chan struct{})
+	defer close(cancelWatcherDone)
+	var cancelOnce sync.Once
+	sendCancel := func() {
+		cancelOnce.Do(func() {
+			cancelCtx, cancelPrompt := context.WithTimeout(context.Background(), 5*time.Second)
+			_ = s.cancelInternalDebugChatPrompt(cancelCtx, spritz, sessionID)
+			cancelPrompt()
+		})
+	}
 
 	dialCtx, dialCancel := context.WithTimeout(runCtx, s.acp.bootstrapDialTimeout)
 	defer dialCancel()
@@ -221,11 +232,48 @@ func (s *server) runInternalDebugChatPrompt(ctx context.Context, spritz *spritzv
 	defer func() {
 		_ = client.close()
 	}()
+	go func() {
+		select {
+		case <-runCtx.Done():
+			select {
+			case <-cancelWatcherDone:
+				return
+			default:
+			}
+			sendCancel()
+		case <-cancelWatcherDone:
+		}
+	}()
 
 	if _, err := client.initialize(runCtx, s.acp.clientInfo, s.acp.clientCapabilities); err != nil {
 		return nil, err
 	}
-	return client.prompt(runCtx, sessionID, message, s.acp.promptSettleTimeout)
+	result, err := client.prompt(runCtx, sessionID, message, s.acp.promptSettleTimeout)
+	if err != nil && runCtx.Err() != nil {
+		sendCancel()
+	}
+	return result, err
+}
+
+func (s *server) cancelInternalDebugChatPrompt(ctx context.Context, spritz *spritzv1.Spritz, sessionID string) error {
+	if spritz == nil || strings.TrimSpace(sessionID) == "" {
+		return nil
+	}
+	dialCtx, cancel := context.WithTimeout(ctx, s.acp.bootstrapDialTimeout)
+	defer cancel()
+
+	instanceConn, _, err := websocket.DefaultDialer.DialContext(dialCtx, s.acpInstanceURL(spritz.Namespace, spritz.Name), nil)
+	if err != nil {
+		return err
+	}
+	client := &acpBootstrapInstanceClient{conn: instanceConn}
+	defer func() {
+		_ = client.close()
+	}()
+	if _, err := client.initialize(ctx, s.acp.clientInfo, s.acp.clientCapabilities); err != nil {
+		return err
+	}
+	return client.cancelPrompt(ctx, sessionID)
 }
 
 func (s *server) auditInternalDebugChat(actorID string, conversation *spritzv1.SpritzConversation, reason, message string, result *acpPromptResult) {
diff --git a/api/internal_debug_chat_test.go b/api/internal_debug_chat_test.go
@@ -21,6 +21,7 @@ type fakeACPDebugChatServerOptions struct {
 	SessionID       string
 	LoadReplay      []map[string]any
 	LoadReplayAfter []map[string]any
+	HoldPrompt      bool
 	PromptChunks    []string
 	PromptError     *acpBootstrapJSONRPCError
 	StopReason      string
@@ -34,6 +35,7 @@ type fakeACPDebugChatServer struct {
 
 	initCalls        int
 	newCalls         int
+	cancelSessionIDs []string
 	loadSessionIDs   []string
 	promptSessionIDs []string
 	promptTexts      []string
@@ -162,6 +164,9 @@ func newFakeACPDebugChatServer(t *testing.T, options fakeACPDebugChatServerOptio
 				fakeServer.promptSessionIDs = append(fakeServer.promptSessionIDs, params.SessionID)
 				fakeServer.promptTexts = append(fakeServer.promptTexts, text)
 				fakeServer.mu.Unlock()
+				if options.HoldPrompt {
+					continue
+				}
 				if options.PromptError != nil {
 					if err := conn.WriteJSON(map[string]any{
 						"jsonrpc": "2.0",
@@ -202,6 +207,16 @@ func newFakeACPDebugChatServer(t *testing.T, options fakeACPDebugChatServerOptio
 				}); err != nil {
 					t.Fatalf("failed to write prompt result: %v", err)
 				}
+			case "session/cancel":
+				var params struct {
+					SessionID string `json:"sessionId"`
+				}
+				if err := json.Unmarshal(message.Params, &params); err != nil {
+					t.Fatalf("failed to decode cancel params: %v", err)
+				}
+				fakeServer.mu.Lock()
+				fakeServer.cancelSessionIDs = append(fakeServer.cancelSessionIDs, params.SessionID)
+				fakeServer.mu.Unlock()
 			default:
 				t.Fatalf("unexpected ACP method %q", message.Method)
 			}
@@ -458,6 +473,44 @@ func TestCleanupInternalDebugConversationUsesFallbackContextWhenCanceled(t *test
 	}
 }
 
+func TestInternalDebugChatSendCancelsPromptWhenPromptTimesOut(t *testing.T) {
+	spritz := readyACPSpritz("tidy-otter", "user-1")
+	fakeACP := newFakeACPDebugChatServer(t, fakeACPDebugChatServerOptions{
+		SessionID:  "session-fresh",
+		HoldPrompt: true,
+	})
+
+	s := newACPTestServer(t, spritz)
+	s.internalAuth = internalAuthConfig{enabled: true, token: "internal-token"}
+	s.acp.instanceURL = func(namespace, name string) string { return fakeACP.url }
+	s.acp.promptTimeout = 20 * time.Millisecond
+
+	e := echo.New()
+	internal := e.Group("", s.internalAuthMiddleware(), s.authMiddleware())
+	internal.POST("/api/internal/v1/debug/chat/send", s.sendInternalDebugChat)
+
+	req := httptest.NewRequest(http.MethodPost, "/api/internal/v1/debug/chat/send", strings.NewReader(`{
+		"target":{"spritzName":"tidy-otter","cwd":"/workspace/app","title":"Debug Run"},
+		"message":"hello from cli"
+	}`))
+	req.Header.Set(internalTokenHeader, "internal-token")
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("X-Spritz-User-Id", "user-1")
+	rec := httptest.NewRecorder()
+	e.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusBadGateway {
+		t.Fatalf("expected status 502, got %d: %s", rec.Code, rec.Body.String())
+	}
+	time.Sleep(20 * time.Millisecond)
+
+	fakeACP.mu.Lock()
+	defer fakeACP.mu.Unlock()
+	if len(fakeACP.cancelSessionIDs) != 1 || fakeACP.cancelSessionIDs[0] != "session-fresh" {
+		t.Fatalf("expected one session/cancel for session-fresh, got %#v", fakeACP.cancelSessionIDs)
+	}
+}
+
 func TestDebugChatRouteIsUnavailableWithoutAuthOrInternalAuth(t *testing.T) {
 	s := newACPTestServer(t)
 	s.auth = authConfig{mode: authModeNone}
diff --git a/cli/src/index.ts b/cli/src/index.ts
@@ -481,13 +481,13 @@ function chatSendUsage() {
   console.log(`Spritz chat send
 
 Usage:
-  spritz chat send (--instance <name> | --conversation <id>) --message <text> [--owner-id <id>] [--reason <text>] [--cwd <path>] [--title <title>] [--namespace <ns>] [--json]
+  spritz chat send (--instance <name> | --conversation <id>) --message <text> [--reason <text>] [--cwd <path>] [--title <title>] [--namespace <ns>] [--json]
 
 Environment:
   SPRITZ_API_URL (default: ${process.env.SPRITZ_API_URL || defaultApiBase})
   SPRITZ_INTERNAL_TOKEN
   SPRITZ_INTERNAL_REQUEST_TIMEOUT_MS
-  SPRITZ_OWNER_ID, SPRITZ_USER_ID, SPRITZ_PROFILE
+  SPRITZ_USER_ID, SPRITZ_PROFILE
 
 Notes:
   --instance creates a new owner-scoped conversation before sending the prompt.
@@ -507,7 +507,7 @@ Usage:
   spritz open <name> [--namespace <ns>]
   spritz terminal <name> [--namespace <ns>] [--session <name>] [--transport <ws|ssh>] [--print]
   spritz ssh <name> [--namespace <ns>] [--session <name>] [--transport <ws|ssh>] [--print]
-  spritz chat send (--instance <name> | --conversation <id>) --message <text> [--owner-id <id>] [--reason <text>] [--cwd <path>] [--title <title>] [--namespace <ns>] [--json]
+  spritz chat send (--instance <name> | --conversation <id>) --message <text> [--reason <text>] [--cwd <path>] [--title <title>] [--namespace <ns>] [--json]
   spritz profile list
   spritz profile current
   spritz profile show [name]
@@ -1468,22 +1468,16 @@ async function main() {
     if (conversationId && (argValue('--cwd') || argValue('--title'))) {
       throw new Error('--cwd and --title are only supported with --instance');
     }
+    if (argValue('--owner-id')?.trim()) {
+      throw new Error('--owner-id is not supported for chat send; authenticate as the caller instead');
+    }
     const message = argValue('--message');
     if (!message?.trim()) {
       throw new Error('--message is required');
     }
-    const { profile } = await resolveProfile({ allowFlag: true });
-    const bearerToken = argValue('--token') || process.env.SPRITZ_BEARER_TOKEN || profile?.bearerToken;
-    const ownerId = argValue('--owner-id')?.trim() || (!bearerToken?.trim() ? await resolveDefaultOwnerId() : undefined);
-    if (!bearerToken?.trim() && !ownerId) {
-      throw new Error('owner id is required when bearer auth is unavailable; use --owner-id or set SPRITZ_OWNER_ID / SPRITZ_USER_ID');
-    }
     const ns = await resolveNamespace();
     const reason = argValue('--reason')?.trim() || 'spz chat send';
     const requestHeaders: Record<string, string> = { 'Content-Type': 'application/json' };
-    if (!bearerToken?.trim() && ownerId) {
-      requestHeaders[headerId] = ownerId;
-    }
     const data = await internalRequest('/internal/v1/debug/chat/send', {
       method: 'POST',
       headers: requestHeaders,
diff --git a/cli/test/chat-send.test.ts b/cli/test/chat-send.test.ts
@@ -56,8 +56,6 @@ test('chat send uses the internal token and prints assistant text by default', a
       'tidy-otter',
       '--message',
       '  hello from cli  ',
-      '--owner-id',
-      'user-123',
       '--cwd',
       '/workspace/app',
       '--title',
@@ -70,6 +68,7 @@ test('chat send uses the internal token and prints assistant text by default', a
         ...process.env,
         SPRITZ_API_URL: `http://127.0.0.1:${address.port}/api`,
         SPRITZ_INTERNAL_TOKEN: 'internal-token',
+        SPRITZ_USER_ID: 'user-123',
         SPRITZ_CONFIG_DIR: mkdtempSync(path.join(os.tmpdir(), 'spz-config-')),
       },
       stdio: ['ignore', 'pipe', 'pipe'],
@@ -144,15 +143,14 @@ test('chat send supports existing conversations and json output', async (t) => {
       'tidy-otter-conv',
       '--message',
       'follow up',
-      '--owner-id',
-      'user-123',
       '--json',
     ],
     {
       env: {
         ...process.env,
         SPRITZ_API_URL: `http://127.0.0.1:${address.port}/api`,
         SPRITZ_INTERNAL_TOKEN: 'internal-token',
+        SPRITZ_USER_ID: 'user-123',
         SPRITZ_CONFIG_DIR: mkdtempSync(path.join(os.tmpdir(), 'spz-config-')),
       },
       stdio: ['ignore', 'pipe', 'pipe'],
@@ -231,8 +229,6 @@ test('chat send does not inject owner header when bearer auth comes from the act
       'tidy-otter',
       '--message',
       'hello from cli',
-      '--owner-id',
-      'victim-id',
     ],
     {
       env: {
@@ -255,6 +251,42 @@ test('chat send does not inject owner header when bearer auth comes from the act
   assert.equal(requestHeaders?.['x-spritz-user-id'], undefined);
 });
 
+test('chat send rejects --owner-id as an auth source', async () => {
+  const child = spawn(
+    process.execPath,
+    [
+      '--import',
+      'tsx',
+      cliPath,
+      'chat',
+      'send',
+      '--instance',
+      'tidy-otter',
+      '--message',
+      'hello from cli',
+      '--owner-id',
+      'victim-id',
+    ],
+    {
+      env: {
+        ...process.env,
+        SPRITZ_INTERNAL_TOKEN: 'internal-token',
+        SPRITZ_CONFIG_DIR: mkdtempSync(path.join(os.tmpdir(), 'spz-config-')),
+      },
+      stdio: ['ignore', 'pipe', 'pipe'],
+    },
+  );
+
+  let stderr = '';
+  child.stderr.on('data', (chunk) => {
+    stderr += chunk.toString();
+  });
+
+  const exitCode = await new Promise<number | null>((resolve) => child.on('exit', resolve));
+  assert.notEqual(exitCode, 0, 'spz chat send should reject --owner-id');
+  assert.match(stderr, /--owner-id is not supported for chat send/);
+});
+
 test('chat send does not require an owner id when bearer auth is available', async (t) => {
   let requestHeaders: http.IncomingHttpHeaders | null = null;
 
