diff --git a/src/browser/components/WorkspaceHeader.tsx b/src/browser/components/WorkspaceHeader.tsx index 5dce2aec82..eec7d21122 100644 --- a/src/browser/components/WorkspaceHeader.tsx +++ b/src/browser/components/WorkspaceHeader.tsx @@ -1,5 +1,6 @@ import React, { useCallback, useEffect, useState } from "react"; import { Pencil } from "lucide-react"; +import { WorkspaceStatusIndicator } from "./WorkspaceStatusIndicator"; import { GitStatusIndicator } from "./GitStatusIndicator"; import { RuntimeBadge } from "./RuntimeBadge"; import { Tooltip, TooltipTrigger, TooltipContent } from "./ui/tooltip"; @@ -79,7 +80,11 @@ export const WorkspaceHeader: React.FC = ({ {namedWorkspacePath} -
+
+
+ +
+ {editorError && {editorError}} diff --git a/src/browser/components/tools/StatusSetToolCall.tsx b/src/browser/components/tools/StatusSetToolCall.tsx index 4643b64ea2..59e4f1a58e 100644 --- a/src/browser/components/tools/StatusSetToolCall.tsx +++ b/src/browser/components/tools/StatusSetToolCall.tsx @@ -22,11 +22,16 @@ export const StatusSetToolCall: React.FC = ({ ? String(result.error) : undefined; + const iconEmoji = "πŸ“‘"; + + const pollLabel = args.poll_interval_s === undefined ? "once" : `${args.poll_interval_s}s`; + const summary = `poll=${pollLabel}: ${args.script.split(/\r?\n/)[0] ?? ""}`; + return ( - - {args.message} + + {summary} {errorMessage && ({errorMessage})} {statusDisplay} diff --git a/src/browser/stores/WorkspaceStore.test.ts b/src/browser/stores/WorkspaceStore.test.ts index d38891fc5d..7e64761698 100644 --- a/src/browser/stores/WorkspaceStore.test.ts +++ b/src/browser/stores/WorkspaceStore.test.ts @@ -219,6 +219,34 @@ describe("WorkspaceStore", () => { }); }); + describe("agent status updates", () => { + it("should surface agent-status-update events in workspace sidebar state", async () => { + const workspaceId = "status-workspace"; + + mockOnChat.mockImplementation(async function* (): AsyncGenerator< + WorkspaceChatMessage, + void, + unknown + > { + await Promise.resolve(); + yield { type: "caught-up" }; + yield { + type: "agent-status-update", + workspaceId, + status: { message: "Building", url: "https://example.com/pr/1" }, + }; + }); + + createAndAddWorkspace(store, workspaceId); + + // Wait for the async iterator loop to process yields + await new Promise((resolve) => setTimeout(resolve, 10)); + + const state = store.getWorkspaceSidebarState(workspaceId); + expect(state.agentStatus).toEqual({ message: "Building", url: "https://example.com/pr/1" }); + }); + }); + describe("syncWorkspaces", () => { it("should add new workspaces", () => { const metadata1: FrontendWorkspaceMetadata = { diff --git a/src/browser/stores/WorkspaceStore.ts b/src/browser/stores/WorkspaceStore.ts index 5c4ffde176..86792a6fbd 100644 --- a/src/browser/stores/WorkspaceStore.ts +++ b/src/browser/stores/WorkspaceStore.ts @@ -43,7 +43,7 @@ export interface WorkspaceState { currentModel: string | null; recencyTimestamp: number | null; todos: TodoItem[]; - agentStatus: { emoji: string; message: string; url?: string } | undefined; + agentStatus: { emoji?: string; message: string; url?: string } | undefined; pendingStreamStartTime: number | null; } @@ -55,7 +55,7 @@ export interface WorkspaceSidebarState { canInterrupt: boolean; currentModel: string | null; recencyTimestamp: number | null; - agentStatus: { emoji: string; message: string; url?: string } | undefined; + agentStatus: { emoji?: string; message: string; url?: string } | undefined; } /** @@ -262,6 +262,15 @@ export class WorkspaceStore { aggregator.handleUsageDelta(data as never); this.usageStore.bump(workspaceId); }, + "agent-status-update": (workspaceId, aggregator, data) => { + // Agent status updates are delivered as chat events (not persisted in history). + // They must: + // 1) be buffered during initial history replay (so they aren't dropped), and + // 2) bump UI state so sidebar/header indicators update immediately. + aggregator.handleMessage(data); + this.states.bump(workspaceId); + }, + "init-start": (workspaceId, aggregator, data) => { aggregator.handleMessage(data); this.states.bump(workspaceId); @@ -943,7 +952,7 @@ export class WorkspaceStore { /** * Check if data is a buffered event type by checking the handler map. - * This ensures isStreamEvent() and processStreamEvent() can never fall out of sync. + * This ensures buffering and processStreamEvent() can never fall out of sync. */ private isBufferedEvent(data: WorkspaceChatMessage): boolean { return "type" in data && data.type in this.bufferedEventHandlers; diff --git a/src/browser/stories/App.chat.stories.tsx b/src/browser/stories/App.chat.stories.tsx index 0fac0723e4..f8995974bc 100644 --- a/src/browser/stories/App.chat.stories.tsx +++ b/src/browser/stories/App.chat.stories.tsx @@ -203,9 +203,8 @@ export const WithAgentStatus: AppStory = { toolCalls: [ createStatusTool( "call-1", - "πŸš€", - "PR #1234 waiting for CI", - "https://github.com/example/repo/pull/1234" + "echo 'πŸš€ PR #1234 waiting for CI https://github.com/example/repo/pull/1234'", + 5 ), ], } diff --git a/src/browser/stories/App.demo.stories.tsx b/src/browser/stories/App.demo.stories.tsx index 949eb34ceb..f1217b82a0 100644 --- a/src/browser/stories/App.demo.stories.tsx +++ b/src/browser/stories/App.demo.stories.tsx @@ -197,9 +197,8 @@ export const Comprehensive: AppStory = { toolCalls: [ createStatusTool( "call-4", - "πŸš€", - "PR #1234 waiting for CI", - "https://github.com/example/repo/pull/1234" + "echo 'πŸš€ PR #1234 waiting for CI https://github.com/example/repo/pull/1234'", + 5 ), ], }), diff --git a/src/browser/stories/mockFactory.ts b/src/browser/stories/mockFactory.ts index 811a31062d..467bf6cd41 100644 --- a/src/browser/stories/mockFactory.ts +++ b/src/browser/stories/mockFactory.ts @@ -291,17 +291,16 @@ export function createTerminalTool( export function createStatusTool( toolCallId: string, - emoji: string, - message: string, - url?: string + script: string, + pollIntervalS?: number ): MuxPart { return { type: "dynamic-tool", toolCallId, toolName: "status_set", state: "output-available", - input: { emoji, message, url }, - output: { success: true, emoji, message, url }, + input: { script, ...(pollIntervalS !== undefined ? { poll_interval_s: pollIntervalS } : {}) }, + output: { success: true }, }; } diff --git a/src/browser/utils/messages/StreamingMessageAggregator.status.test.ts b/src/browser/utils/messages/StreamingMessageAggregator.status.test.ts index ac6a4befa6..c3dafe3aa9 100644 --- a/src/browser/utils/messages/StreamingMessageAggregator.status.test.ts +++ b/src/browser/utils/messages/StreamingMessageAggregator.status.test.ts @@ -1,103 +1,52 @@ -import { afterAll, afterEach, beforeEach, describe, expect, it } from "bun:test"; -import { getStatusStateKey } from "@/common/constants/storage"; +import { describe, expect, it } from "bun:test"; import { StreamingMessageAggregator } from "./StreamingMessageAggregator"; -const originalLocalStorage: Storage | undefined = (globalThis as { localStorage?: Storage }) - .localStorage; - -const createMockLocalStorage = () => { - const store = new Map(); - return { - get length() { - return store.size; - }, - key: (index: number) => Array.from(store.keys())[index] ?? null, - getItem: (key: string) => (store.has(key) ? store.get(key)! : null), - setItem: (key: string, value: string) => { - store.set(key, value); - }, - removeItem: (key: string) => { - store.delete(key); - }, - clear: () => { - store.clear(); - }, - } satisfies Storage; -}; - -beforeEach(() => { - const mock = createMockLocalStorage(); - Object.defineProperty(globalThis, "localStorage", { - value: mock, - configurable: true, - }); -}); - -afterEach(() => { - const ls = (globalThis as { localStorage?: Storage }).localStorage; - ls?.clear?.(); -}); - -afterAll(() => { - if (originalLocalStorage !== undefined) { - Object.defineProperty(globalThis, "localStorage", { value: originalLocalStorage }); - } else { - delete (globalThis as { localStorage?: Storage }).localStorage; - } -}); - describe("StreamingMessageAggregator - Agent Status", () => { it("should start with undefined agent status", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); expect(aggregator.getAgentStatus()).toBeUndefined(); }); - it("should update agent status when status_set tool succeeds", () => { + it("should update agent status when receiving agent-status-update", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - const messageId = "msg1"; - const toolCallId = "tool1"; - // Start a stream - aggregator.handleStreamStart({ - type: "stream-start", + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - model: "test-model", - historySequence: 1, + status: { + emoji: "πŸš€", + message: "PR #1 checks running", + url: "https://github.com/example/repo/pull/1", + }, }); - // Add a status_set tool call - aggregator.handleToolCallStart({ - type: "tool-call-start", - workspaceId: "workspace1", - messageId, - toolCallId, - toolName: "status_set", - args: { emoji: "πŸ”", message: "Analyzing code" }, - tokens: 10, - timestamp: Date.now(), + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "πŸš€", + message: "PR #1 checks running", + url: "https://github.com/example/repo/pull/1", }); - // Complete the tool call - aggregator.handleToolCallEnd({ - type: "tool-call-end", + // URL should persist when subsequent updates omit it + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - toolCallId, - toolName: "status_set", - result: { success: true, emoji: "πŸ”", message: "Analyzing code" }, - timestamp: Date.now(), + status: { + emoji: "🟑", + message: "PR #1 mergeable", + }, }); - const status = aggregator.getAgentStatus(); - expect(status).toBeDefined(); - expect(status?.emoji).toBe("πŸ”"); - expect(status?.message).toBe("Analyzing code"); + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "🟑", + message: "PR #1 mergeable", + url: "https://github.com/example/repo/pull/1", + }); }); - it("should update agent status multiple times", () => { + it("should not update agent status from status_set tool results", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); const messageId = "msg1"; + const toolCallId = "tool1"; // Start a stream aggregator.handleStreamStart({ @@ -108,54 +57,32 @@ describe("StreamingMessageAggregator - Agent Status", () => { historySequence: 1, }); - // First status_set - aggregator.handleToolCallStart({ - type: "tool-call-start", - workspaceId: "workspace1", - messageId, - toolCallId: "tool1", - toolName: "status_set", - args: { emoji: "πŸ”", message: "Analyzing" }, - tokens: 10, - timestamp: Date.now(), - }); - - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId, - toolCallId: "tool1", - toolName: "status_set", - result: { success: true, emoji: "πŸ”", message: "Analyzing" }, - timestamp: Date.now(), - }); - - expect(aggregator.getAgentStatus()?.emoji).toBe("πŸ”"); - - // Second status_set + // Add a status_set tool call aggregator.handleToolCallStart({ type: "tool-call-start", workspaceId: "workspace1", messageId, - toolCallId: "tool2", + toolCallId, toolName: "status_set", - args: { emoji: "πŸ“", message: "Writing" }, + args: { + script: "echo 'πŸš€ PR #1 https://github.com/example/repo/pull/1'", + }, tokens: 10, timestamp: Date.now(), }); + // Complete the tool call (success is just an acknowledgement) aggregator.handleToolCallEnd({ type: "tool-call-end", workspaceId: "workspace1", messageId, - toolCallId: "tool2", + toolCallId, toolName: "status_set", - result: { success: true, emoji: "πŸ“", message: "Writing" }, + result: { success: true }, timestamp: Date.now(), }); - expect(aggregator.getAgentStatus()?.emoji).toBe("πŸ“"); - expect(aggregator.getAgentStatus()?.message).toBe("Writing"); + expect(aggregator.getAgentStatus()).toBeUndefined(); }); it("should persist agent status after stream ends", () => { @@ -171,30 +98,17 @@ describe("StreamingMessageAggregator - Agent Status", () => { historySequence: 1, }); - // Set status - aggregator.handleToolCallStart({ - type: "tool-call-start", - workspaceId: "workspace1", - messageId, - toolCallId: "tool1", - toolName: "status_set", - args: { emoji: "πŸ”", message: "Working" }, - tokens: 10, - timestamp: Date.now(), - }); - - aggregator.handleToolCallEnd({ - type: "tool-call-end", + // Status arrives via agent-status-update + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - toolCallId: "tool1", - toolName: "status_set", - result: { success: true, emoji: "πŸ”", message: "Working" }, - timestamp: Date.now(), + status: { + emoji: "πŸ”", + message: "Working", + url: "https://github.com/example/repo/pull/1", + }, }); - expect(aggregator.getAgentStatus()).toBeDefined(); - // End the stream aggregator.handleStreamEnd({ type: "stream-end", @@ -204,16 +118,17 @@ describe("StreamingMessageAggregator - Agent Status", () => { parts: [], }); - // Status should persist after stream ends (unlike todos) - expect(aggregator.getAgentStatus()).toBeDefined(); - expect(aggregator.getAgentStatus()?.emoji).toBe("πŸ”"); + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "πŸ”", + message: "Working", + url: "https://github.com/example/repo/pull/1", + }); }); - it("should not update agent status if tool call fails", () => { + it("should keep agent status unchanged when status_set tool call fails", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); const messageId = "msg1"; - // Start a stream aggregator.handleStreamStart({ type: "stream-start", workspaceId: "workspace1", @@ -222,19 +137,18 @@ describe("StreamingMessageAggregator - Agent Status", () => { historySequence: 1, }); - // Add a status_set tool call - aggregator.handleToolCallStart({ - type: "tool-call-start", + // Establish a status via agent-status-update + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - toolCallId: "tool1", - toolName: "status_set", - args: { emoji: "πŸ”", message: "Analyzing" }, - tokens: 10, - timestamp: Date.now(), + status: { + emoji: "🟑", + message: "CI running", + url: "https://github.com/example/repo/pull/1", + }, }); - // Complete with failure + // A failing status_set tool call should not wipe status aggregator.handleToolCallEnd({ type: "tool-call-end", workspaceId: "workspace1", @@ -245,14 +159,16 @@ describe("StreamingMessageAggregator - Agent Status", () => { timestamp: Date.now(), }); - // Status should remain undefined - expect(aggregator.getAgentStatus()).toBeUndefined(); + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "🟑", + message: "CI running", + url: "https://github.com/example/repo/pull/1", + }); }); - it("should clear agent status when new user message arrives", () => { + it("should not clear agent status when new user message arrives", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - // Start first stream and set status aggregator.handleStreamStart({ type: "stream-start", workspaceId: "workspace1", @@ -261,42 +177,19 @@ describe("StreamingMessageAggregator - Agent Status", () => { historySequence: 1, }); - aggregator.handleToolCallStart({ - type: "tool-call-start", - workspaceId: "workspace1", - messageId: "msg1", - toolCallId: "tool1", - toolName: "status_set", - args: { emoji: "πŸ”", message: "First task" }, - tokens: 10, - timestamp: Date.now(), - }); - - aggregator.handleToolCallEnd({ - type: "tool-call-end", + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId: "msg1", - toolCallId: "tool1", - toolName: "status_set", - result: { success: true, emoji: "πŸ”", message: "First task" }, - timestamp: Date.now(), - }); - - expect(aggregator.getAgentStatus()?.message).toBe("First task"); - - // End first stream - aggregator.handleStreamEnd({ - type: "stream-end", - workspaceId: "workspace1", - messageId: "msg1", - metadata: { model: "test-model" }, - parts: [], + status: { + emoji: "πŸ”", + message: "First task", + url: "https://github.com/example/repo/pull/1", + }, }); - // Status persists after stream ends expect(aggregator.getAgentStatus()?.message).toBe("First task"); - // User sends a NEW message - status should be cleared + // User sends a NEW message - status should persist const newUserMessage = { type: "message" as const, id: "msg2", @@ -306,8 +199,7 @@ describe("StreamingMessageAggregator - Agent Status", () => { }; aggregator.handleMessage(newUserMessage); - // Status should be cleared on new user message - expect(aggregator.getAgentStatus()).toBeUndefined(); + expect(aggregator.getAgentStatus()?.url).toBe("https://github.com/example/repo/pull/1"); }); it("should show 'failed' status in UI when status_set validation fails", () => { @@ -323,14 +215,14 @@ describe("StreamingMessageAggregator - Agent Status", () => { historySequence: 1, }); - // Add a status_set tool call with invalid emoji + // Add a status_set tool call aggregator.handleToolCallStart({ type: "tool-call-start", workspaceId: "workspace1", messageId, toolCallId: "tool1", toolName: "status_set", - args: { emoji: "not-an-emoji", message: "test" }, + args: { script: "echo 'not important'" }, tokens: 10, timestamp: Date.now(), }); @@ -369,7 +261,7 @@ describe("StreamingMessageAggregator - Agent Status", () => { expect(aggregator.getAgentStatus()).toBeUndefined(); }); - it("should show 'completed' status in UI when status_set validation succeeds", () => { + it("should show 'completed' status in UI when status_set succeeds", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); const messageId = "msg1"; @@ -389,7 +281,9 @@ describe("StreamingMessageAggregator - Agent Status", () => { messageId, toolCallId: "tool1", toolName: "status_set", - args: { emoji: "πŸ”", message: "Analyzing code" }, + args: { + script: "echo 'πŸš€ PR #1 https://github.com/example/repo/pull/1'", + }, tokens: 10, timestamp: Date.now(), }); @@ -401,7 +295,7 @@ describe("StreamingMessageAggregator - Agent Status", () => { messageId, toolCallId: "tool1", toolName: "status_set", - result: { success: true, emoji: "πŸ”", message: "Analyzing code" }, + result: { success: true }, timestamp: Date.now(), }); @@ -424,14 +318,11 @@ describe("StreamingMessageAggregator - Agent Status", () => { expect(toolMessage.toolName).toBe("status_set"); } - // And status SHOULD be updated in aggregator - const status = aggregator.getAgentStatus(); - expect(status).toBeDefined(); - expect(status?.emoji).toBe("πŸ”"); - expect(status?.message).toBe("Analyzing code"); + // And agent status is unchanged (status is delivered via agent-status-update events) + expect(aggregator.getAgentStatus()).toBeUndefined(); }); - it("should reconstruct agentStatus when loading historical messages", () => { + it("should not reconstruct agent status from historical status_set tool calls", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); // Create historical messages with a completed status_set tool call @@ -452,8 +343,8 @@ describe("StreamingMessageAggregator - Agent Status", () => { toolCallId: "tool1", toolName: "status_set", state: "output-available" as const, - input: { emoji: "πŸ”", message: "Analyzing code" }, - output: { success: true, emoji: "πŸ”", message: "Analyzing code" }, + input: { script: "echo 'Analyzing code'" }, + output: { success: true }, timestamp: Date.now(), }, ], @@ -464,491 +355,114 @@ describe("StreamingMessageAggregator - Agent Status", () => { // Load historical messages aggregator.loadHistoricalMessages(historicalMessages); - // Status should be reconstructed from the historical tool call - const status = aggregator.getAgentStatus(); - expect(status).toBeDefined(); - expect(status?.emoji).toBe("πŸ”"); - expect(status?.message).toBe("Analyzing code"); - }); - - it("should use most recent status_set when loading multiple historical messages", () => { - const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - - // Create historical messages with multiple status_set calls - const historicalMessages = [ - { - id: "msg1", - role: "assistant" as const, - parts: [ - { - type: "dynamic-tool" as const, - toolCallId: "tool1", - toolName: "status_set", - state: "output-available" as const, - input: { emoji: "πŸ”", message: "First status" }, - output: { success: true, emoji: "πŸ”", message: "First status" }, - timestamp: Date.now(), - }, - ], - metadata: { timestamp: Date.now(), historySequence: 1 }, - }, - { - id: "msg2", - role: "assistant" as const, - parts: [ - { - type: "dynamic-tool" as const, - toolCallId: "tool2", - toolName: "status_set", - state: "output-available" as const, - input: { emoji: "πŸ“", message: "Second status" }, - output: { success: true, emoji: "πŸ“", message: "Second status" }, - timestamp: Date.now(), - }, - ], - metadata: { timestamp: Date.now(), historySequence: 2 }, - }, - ]; - - // Load historical messages - aggregator.loadHistoricalMessages(historicalMessages); - - // Should use the most recent (last processed) status - const status = aggregator.getAgentStatus(); - expect(status?.emoji).toBe("πŸ“"); - expect(status?.message).toBe("Second status"); - }); - - it("should not reconstruct status from failed status_set in historical messages", () => { - const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - - // Create historical message with failed status_set - const historicalMessages = [ - { - id: "msg1", - role: "assistant" as const, - parts: [ - { - type: "dynamic-tool" as const, - toolCallId: "tool1", - toolName: "status_set", - state: "output-available" as const, - input: { emoji: "not-emoji", message: "test" }, - output: { success: false, error: "emoji must be a single emoji character" }, - timestamp: Date.now(), - }, - ], - metadata: { timestamp: Date.now(), historySequence: 1 }, - }, - ]; - - // Load historical messages - aggregator.loadHistoricalMessages(historicalMessages); - - // Status should remain undefined (failed validation) + // status_set does not reconstruct agent status from history (status is ephemeral and persisted separately) expect(aggregator.getAgentStatus()).toBeUndefined(); }); - it("should retain last status_set even if later assistant messages omit it", () => { + it("should store URL when provided in agent-status-update", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - const historicalMessages = [ - { - id: "assistant1", - role: "assistant" as const, - parts: [ - { - type: "dynamic-tool" as const, - toolCallId: "tool1", - toolName: "status_set", - state: "output-available" as const, - input: { emoji: "πŸ§ͺ", message: "Running tests" }, - output: { success: true, emoji: "πŸ§ͺ", message: "Running tests" }, - timestamp: 1000, - }, - ], - metadata: { timestamp: 1000, historySequence: 1 }, - }, - { - id: "assistant2", - role: "assistant" as const, - parts: [{ type: "text" as const, text: "[compaction summary]" }], - metadata: { timestamp: 2000, historySequence: 2 }, - }, - ]; - - aggregator.loadHistoricalMessages(historicalMessages); - - const status = aggregator.getAgentStatus(); - expect(status?.emoji).toBe("πŸ§ͺ"); - expect(status?.message).toBe("Running tests"); - }); - - it("should restore persisted status when history is compacted away", () => { - const workspaceId = "workspace1"; - const persistedStatus = { - emoji: "πŸ”—", - message: "PR open", - url: "https://example.com/pr/123", - } as const; - localStorage.setItem(getStatusStateKey(workspaceId), JSON.stringify(persistedStatus)); - - const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z", workspaceId); - - // History with no status_set (e.g., after compaction removes older tool calls) - const historicalMessages = [ - { - id: "assistant2", - role: "assistant" as const, - parts: [{ type: "text" as const, text: "[compacted history]" }], - metadata: { timestamp: 3000, historySequence: 1 }, - }, - ]; - - aggregator.loadHistoricalMessages(historicalMessages); - - expect(aggregator.getAgentStatus()).toEqual(persistedStatus); - }); - - it("should use truncated message from output, not original input", () => { - const aggregator = new StreamingMessageAggregator(new Date().toISOString()); - - const messageId = "msg1"; - const toolCallId = "tool1"; - - // Start stream - aggregator.handleStreamStart({ - type: "stream-start", - workspaceId: "workspace1", - messageId, - model: "test-model", - historySequence: 1, - }); - - // Status_set with long message (would be truncated by backend) - const longMessage = "a".repeat(100); // 100 chars, exceeds 60 char limit - const truncatedMessage = "a".repeat(59) + "…"; // What backend returns - - aggregator.handleToolCallStart({ - type: "tool-call-start", - workspaceId: "workspace1", - messageId, - toolCallId, - toolName: "status_set", - args: { emoji: "βœ…", message: longMessage }, - tokens: 10, - timestamp: Date.now(), - }); - - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId, - toolCallId, - toolName: "status_set", - result: { success: true, emoji: "βœ…", message: truncatedMessage }, - timestamp: Date.now(), - }); - - // Should use truncated message from output, not the original input - const status = aggregator.getAgentStatus(); - expect(status).toEqual({ emoji: "βœ…", message: truncatedMessage }); - expect(status?.message.length).toBe(60); - }); - - it("should store URL when provided in status_set", () => { - const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - const messageId = "msg1"; - const toolCallId = "tool1"; - - // Start a stream - aggregator.handleStreamStart({ - type: "stream-start", - workspaceId: "workspace1", - messageId, - model: "test-model", - historySequence: 1, - }); - - // Add a status_set tool call with URL const testUrl = "https://github.com/owner/repo/pull/123"; - aggregator.handleToolCallStart({ - type: "tool-call-start", + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - toolCallId, - toolName: "status_set", - args: { emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - tokens: 10, - timestamp: Date.now(), + status: { + emoji: "πŸ”—", + message: "PR submitted", + url: testUrl, + }, }); - // Complete the tool call - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId, - toolCallId, - toolName: "status_set", - result: { success: true, emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - timestamp: Date.now(), + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "πŸ”—", + message: "PR submitted", + url: testUrl, }); - - const status = aggregator.getAgentStatus(); - expect(status).toBeDefined(); - expect(status?.emoji).toBe("πŸ”—"); - expect(status?.message).toBe("PR submitted"); - expect(status?.url).toBe(testUrl); }); - it("should persist URL across status updates until explicitly replaced", () => { + it("should persist URL across agent-status-update events until explicitly replaced", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - const messageId = "msg1"; - // Start a stream - aggregator.handleStreamStart({ - type: "stream-start", - workspaceId: "workspace1", - messageId, - model: "test-model", - historySequence: 1, - }); - - // First status with URL const testUrl = "https://github.com/owner/repo/pull/123"; - aggregator.handleToolCallStart({ - type: "tool-call-start", + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - toolCallId: "tool1", - toolName: "status_set", - args: { emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - tokens: 10, - timestamp: Date.now(), - }); - - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId, - toolCallId: "tool1", - toolName: "status_set", - result: { success: true, emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - timestamp: Date.now(), + status: { emoji: "πŸ”—", message: "PR submitted", url: testUrl }, }); expect(aggregator.getAgentStatus()?.url).toBe(testUrl); - // Second status without URL - should keep previous URL - aggregator.handleToolCallStart({ - type: "tool-call-start", + // Update without URL - should keep previous URL + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - toolCallId: "tool2", - toolName: "status_set", - args: { emoji: "βœ…", message: "Done" }, - tokens: 10, - timestamp: Date.now(), + status: { emoji: "βœ…", message: "Done" }, }); - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId, - toolCallId: "tool2", - toolName: "status_set", - result: { success: true, emoji: "βœ…", message: "Done" }, - timestamp: Date.now(), + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "βœ…", + message: "Done", + url: testUrl, }); - const statusAfterUpdate = aggregator.getAgentStatus(); - expect(statusAfterUpdate?.emoji).toBe("βœ…"); - expect(statusAfterUpdate?.message).toBe("Done"); - expect(statusAfterUpdate?.url).toBe(testUrl); // URL persists - - // Third status with different URL - should replace + // Update with a different URL - should replace const newUrl = "https://github.com/owner/repo/pull/456"; - aggregator.handleToolCallStart({ - type: "tool-call-start", + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId, - toolCallId: "tool3", - toolName: "status_set", - args: { emoji: "πŸ”„", message: "New PR", url: newUrl }, - tokens: 10, - timestamp: Date.now(), + status: { emoji: "πŸ”„", message: "New PR", url: newUrl }, }); - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId, - toolCallId: "tool3", - toolName: "status_set", - result: { success: true, emoji: "πŸ”„", message: "New PR", url: newUrl }, - timestamp: Date.now(), + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "πŸ”„", + message: "New PR", + url: newUrl, }); - - const finalStatus = aggregator.getAgentStatus(); - expect(finalStatus?.emoji).toBe("πŸ”„"); - expect(finalStatus?.message).toBe("New PR"); - expect(finalStatus?.url).toBe(newUrl); // URL replaced }); - it("should persist URL even after status is cleared by new stream start", () => { + it("should persist URL across user turns and stream boundaries", () => { const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - const messageId1 = "msg1"; - // Start first stream - aggregator.handleStreamStart({ - type: "stream-start", - workspaceId: "workspace1", - messageId: messageId1, - model: "test-model", - historySequence: 1, - }); - - // Set status with URL in first stream const testUrl = "https://github.com/owner/repo/pull/123"; - aggregator.handleToolCallStart({ - type: "tool-call-start", + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId: messageId1, - toolCallId: "tool1", - toolName: "status_set", - args: { emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - tokens: 10, - timestamp: Date.now(), - }); - - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId: messageId1, - toolCallId: "tool1", - toolName: "status_set", - result: { success: true, emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - timestamp: Date.now(), + status: { emoji: "πŸ”—", message: "PR submitted", url: testUrl }, }); expect(aggregator.getAgentStatus()?.url).toBe(testUrl); - // User sends a new message, which clears the status - const userMessage = { + // User sends a follow-up + aggregator.handleMessage({ type: "message" as const, id: "user1", role: "user" as const, parts: [{ type: "text" as const, text: "Continue" }], metadata: { timestamp: Date.now(), historySequence: 2 }, - }; - aggregator.handleMessage(userMessage); + }); - expect(aggregator.getAgentStatus()).toBeUndefined(); // Status cleared + expect(aggregator.getAgentStatus()?.url).toBe(testUrl); - // Start second stream - const messageId2 = "msg2"; + // New stream starts aggregator.handleStreamStart({ type: "stream-start", workspaceId: "workspace1", - messageId: messageId2, + messageId: "msg2", model: "test-model", historySequence: 2, }); - // Set new status WITHOUT URL - should use the last URL ever seen - aggregator.handleToolCallStart({ - type: "tool-call-start", + // Status update without URL retains last URL + aggregator.handleMessage({ + type: "agent-status-update", workspaceId: "workspace1", - messageId: messageId2, - toolCallId: "tool2", - toolName: "status_set", - args: { emoji: "βœ…", message: "Tests passed" }, - tokens: 10, - timestamp: Date.now(), + status: { emoji: "βœ…", message: "Tests passed" }, }); - aggregator.handleToolCallEnd({ - type: "tool-call-end", - workspaceId: "workspace1", - messageId: messageId2, - toolCallId: "tool2", - toolName: "status_set", - result: { success: true, emoji: "βœ…", message: "Tests passed" }, - timestamp: Date.now(), + expect(aggregator.getAgentStatus()).toEqual({ + emoji: "βœ…", + message: "Tests passed", + url: testUrl, }); - - const finalStatus = aggregator.getAgentStatus(); - expect(finalStatus?.emoji).toBe("βœ…"); - expect(finalStatus?.message).toBe("Tests passed"); - expect(finalStatus?.url).toBe(testUrl); // URL from previous stream persists! - }); - - it("should persist URL across multiple assistant messages when loading from history", () => { - // Regression test: URL should persist even when only the most recent assistant message - // has a status_set without a URL - the URL from an earlier message should be used - const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z"); - const testUrl = "https://github.com/owner/repo/pull/123"; - - // Historical messages: first assistant sets URL, second assistant updates status without URL - const historicalMessages = [ - { - id: "user1", - role: "user" as const, - parts: [{ type: "text" as const, text: "Make a PR" }], - metadata: { timestamp: 1000, historySequence: 1 }, - }, - { - id: "assistant1", - role: "assistant" as const, - parts: [ - { - type: "dynamic-tool" as const, - toolName: "status_set", - toolCallId: "tool1", - state: "output-available" as const, - input: { emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - output: { success: true, emoji: "πŸ”—", message: "PR submitted", url: testUrl }, - timestamp: 1001, - tokens: 10, - }, - ], - metadata: { timestamp: 1001, historySequence: 2 }, - }, - { - id: "user2", - role: "user" as const, - parts: [{ type: "text" as const, text: "Continue" }], - metadata: { timestamp: 2000, historySequence: 3 }, - }, - { - id: "assistant2", - role: "assistant" as const, - parts: [ - { - type: "dynamic-tool" as const, - toolName: "status_set", - toolCallId: "tool2", - state: "output-available" as const, - input: { emoji: "βœ…", message: "Tests passed" }, - output: { success: true, emoji: "βœ…", message: "Tests passed" }, // No URL! - timestamp: 2001, - tokens: 10, - }, - ], - metadata: { timestamp: 2001, historySequence: 4 }, - }, - ]; - - aggregator.loadHistoricalMessages(historicalMessages); - - const status = aggregator.getAgentStatus(); - expect(status?.emoji).toBe("βœ…"); - expect(status?.message).toBe("Tests passed"); - // URL from the first assistant message should persist! - expect(status?.url).toBe(testUrl); }); - - // Note: URL persistence through compaction is handled via localStorage, - // which is tested in integration tests. The aggregator saves lastStatusUrl - // to localStorage when it changes, and loads it on construction. }); diff --git a/src/browser/utils/messages/StreamingMessageAggregator.test.ts b/src/browser/utils/messages/StreamingMessageAggregator.test.ts index 6f4e6e090a..b9676565a4 100644 --- a/src/browser/utils/messages/StreamingMessageAggregator.test.ts +++ b/src/browser/utils/messages/StreamingMessageAggregator.test.ts @@ -289,7 +289,7 @@ describe("StreamingMessageAggregator", () => { expect(aggregator2.getCurrentTodos()).toHaveLength(0); }); - test("should reconstruct agentStatus but NOT todos when no active stream", () => { + test("should not reconstruct agentStatus but NOT todos when no active stream", () => { const aggregator = new StreamingMessageAggregator(TEST_CREATED_AT); const historicalMessage = { @@ -312,8 +312,8 @@ describe("StreamingMessageAggregator", () => { toolCallId: "tool2", toolName: "status_set", state: "output-available" as const, - input: { emoji: "πŸ”§", message: "Working on it" }, - output: { success: true, emoji: "πŸ”§", message: "Working on it" }, + input: { script: "echo 'Working on it'" }, + output: { success: true }, }, ], metadata: { @@ -326,8 +326,8 @@ describe("StreamingMessageAggregator", () => { // Load without active stream aggregator.loadHistoricalMessages([historicalMessage], false); - // agentStatus should be reconstructed (persists across sessions) - expect(aggregator.getAgentStatus()).toEqual({ emoji: "πŸ”§", message: "Working on it" }); + // agentStatus is not reconstructed from status_set tool calls. + expect(aggregator.getAgentStatus()).toBeUndefined(); // TODOs should NOT be reconstructed (stream-scoped) expect(aggregator.getCurrentTodos()).toHaveLength(0); diff --git a/src/browser/utils/messages/StreamingMessageAggregator.ts b/src/browser/utils/messages/StreamingMessageAggregator.ts index 0c8e0909dc..5ae8725048 100644 --- a/src/browser/utils/messages/StreamingMessageAggregator.ts +++ b/src/browser/utils/messages/StreamingMessageAggregator.ts @@ -19,30 +19,32 @@ import type { ReasoningEndEvent, } from "@/common/types/stream"; import type { LanguageModelV2Usage } from "@ai-sdk/provider"; -import type { TodoItem, StatusSetToolResult } from "@/common/types/tools"; +import type { TodoItem } from "@/common/types/tools"; import type { WorkspaceChatMessage, StreamErrorMessage, DeleteMessage } from "@/common/orpc/types"; -import { isInitStart, isInitOutput, isInitEnd, isMuxMessage } from "@/common/orpc/types"; +import { + isInitStart, + isInitOutput, + isInitEnd, + isMuxMessage, + isAgentStatusUpdate, +} from "@/common/orpc/types"; import type { DynamicToolPart, DynamicToolPartPending, DynamicToolPartAvailable, } from "@/common/types/toolParts"; import { isDynamicToolPart } from "@/common/types/toolParts"; -import { z } from "zod"; import { createDeltaStorage, type DeltaRecordStorage } from "./StreamingTPSCalculator"; import { computeRecencyTimestamp } from "./recency"; -import { getStatusStateKey } from "@/common/constants/storage"; // Maximum number of messages to display in the DOM for performance // Full history is still maintained internally for token counting and stats -const AgentStatusSchema = z.object({ - emoji: z.string(), - message: z.string(), - url: z.string().optional(), -}); - -type AgentStatus = z.infer; +interface AgentStatus { + emoji?: string; + message: string; + url?: string; +} const MAX_DISPLAYED_MESSAGES = 128; interface StreamingContext { @@ -110,9 +112,6 @@ export class StreamingMessageAggregator { // Last URL set via status_set - kept in memory to reuse when later calls omit url private lastStatusUrl: string | undefined = undefined; - // Workspace ID for localStorage persistence - private readonly workspaceId: string | undefined; - // Workspace init hook state (ephemeral, not persisted to history) private initState: { status: "running" | "success" | "error"; @@ -133,55 +132,10 @@ export class StreamingMessageAggregator { // REQUIRED: Backend guarantees every workspace has createdAt via config.ts private readonly createdAt: string; - constructor(createdAt: string, workspaceId?: string) { + constructor(createdAt: string, _workspaceId?: string) { this.createdAt = createdAt; - this.workspaceId = workspaceId; - // Load persisted agent status from localStorage - if (workspaceId) { - const persistedStatus = this.loadPersistedAgentStatus(); - if (persistedStatus) { - this.agentStatus = persistedStatus; - this.lastStatusUrl = persistedStatus.url; - } - } this.updateRecency(); } - - /** Load persisted agent status from localStorage */ - private loadPersistedAgentStatus(): AgentStatus | undefined { - if (!this.workspaceId) return undefined; - try { - const stored = localStorage.getItem(getStatusStateKey(this.workspaceId)); - if (!stored) return undefined; - const parsed = AgentStatusSchema.safeParse(JSON.parse(stored)); - return parsed.success ? parsed.data : undefined; - } catch { - // Ignore localStorage errors or JSON parse failures - } - return undefined; - } - - /** Persist agent status to localStorage */ - private savePersistedAgentStatus(status: AgentStatus): void { - if (!this.workspaceId) return; - const parsed = AgentStatusSchema.safeParse(status); - if (!parsed.success) return; - try { - localStorage.setItem(getStatusStateKey(this.workspaceId), JSON.stringify(parsed.data)); - } catch { - // Ignore localStorage errors - } - } - - /** Remove persisted agent status from localStorage */ - private clearPersistedAgentStatus(): void { - if (!this.workspaceId) return; - try { - localStorage.removeItem(getStatusStateKey(this.workspaceId)); - } catch { - // Ignore localStorage errors - } - } private invalidateCache(): void { this.cachedAllMessages = null; this.cachedDisplayedMessages = null; @@ -320,10 +274,9 @@ export class StreamingMessageAggregator { // Replay historical messages in order to reconstruct derived state for (const message of chronologicalMessages) { if (message.role === "user") { - // Mirror live behavior: clear stream-scoped state on new user turn - // but keep persisted status for fallback on reload. + // Mirror live behavior: clear stream-scoped state on new user turn. + // Agent status is NOT stream-scoped and is intentionally not cleared. this.currentTodos = []; - this.agentStatus = undefined; continue; } @@ -336,15 +289,6 @@ export class StreamingMessageAggregator { } } - // If history was compacted away from the last status_set, fall back to persisted status - if (!this.agentStatus) { - const persistedStatus = this.loadPersistedAgentStatus(); - if (persistedStatus) { - this.agentStatus = persistedStatus; - this.lastStatusUrl = persistedStatus.url; - } - } - this.invalidateCache(); } @@ -683,25 +627,8 @@ export class StreamingMessageAggregator { } } - // Update agent status if this was a successful status_set - // agentStatus persists: update both during streaming and on historical reload - // Use output instead of input to get the truncated message - if (toolName === "status_set" && hasSuccessResult(output)) { - const result = output as Extract; - - // Use the provided URL, or fall back to the last URL ever set - const url = result.url ?? this.lastStatusUrl; - if (url) { - this.lastStatusUrl = url; - } - - this.agentStatus = { - emoji: result.emoji, - message: result.message, - url, - }; - this.savePersistedAgentStatus(this.agentStatus); - } + // status_set updates are delivered via "agent-status-update" events. + // Tool results are only acknowledgements ({ success: true }). } handleToolCallEnd(data: ToolCallEndEvent): void { @@ -795,6 +722,21 @@ export class StreamingMessageAggregator { } // Handle regular messages (user messages, historical messages) + if (isAgentStatusUpdate(data)) { + const effectiveUrl = data.status.url ?? this.lastStatusUrl; + if (effectiveUrl) { + this.lastStatusUrl = effectiveUrl; + } + + this.agentStatus = { + ...(data.status.emoji ? { emoji: data.status.emoji } : {}), + message: data.status.message, + ...(effectiveUrl ? { url: effectiveUrl } : {}), + }; + this.invalidateCache(); + return; + } + // Check if it's a MuxMessage (has role property but no type) if (isMuxMessage(data)) { const incomingMessage = data; @@ -830,12 +772,11 @@ export class StreamingMessageAggregator { // If this is a user message, clear derived state and record timestamp if (incomingMessage.role === "user") { - // Clear derived state (todos, agentStatus) for new conversation turn - // This ensures consistent behavior whether loading from history or processing live events - // since stream-start/stream-end events are not persisted in chat.jsonl + // Clear derived state for new conversation turn. + // TODOs are stream-scoped; agent status is intentionally NOT cleared. + // This ensures the user can still see/click important links (e.g., PR URL) + // even after the assistant stops streaming or the user sends a follow-up message. this.currentTodos = []; - this.agentStatus = undefined; - this.clearPersistedAgentStatus(); this.setPendingStreamStartTime(Date.now()); } diff --git a/src/browser/utils/ui/statusTooltip.tsx b/src/browser/utils/ui/statusTooltip.tsx index 7f67bc66a6..e8a70a8c69 100644 --- a/src/browser/utils/ui/statusTooltip.tsx +++ b/src/browser/utils/ui/statusTooltip.tsx @@ -9,7 +9,7 @@ import { formatRelativeTime } from "@/browser/utils/ui/dateTime"; export function getStatusTooltip(options: { isStreaming: boolean; streamingModel: string | null; - agentStatus?: { emoji: string; message: string; url?: string }; + agentStatus?: { emoji?: string; message: string; url?: string }; isUnread?: boolean; recencyTimestamp?: number | null; }): React.ReactNode { diff --git a/src/common/constants/storage.ts b/src/common/constants/storage.ts index beb43ba657..c11869c27c 100644 --- a/src/common/constants/storage.ts +++ b/src/common/constants/storage.ts @@ -225,15 +225,6 @@ export function getFileTreeExpandStateKey(workspaceId: string): string { return `fileTreeExpandState:${workspaceId}`; } -/** - * Get the localStorage key for persisted agent status for a workspace - * Stores the most recent successful status_set payload (emoji, message, url) - * Format: "statusState:{workspaceId}" - */ -export function getStatusStateKey(workspaceId: string): string { - return `statusState:${workspaceId}`; -} - /** * Right sidebar tab selection (global) * Format: "right-sidebar-tab" @@ -297,7 +288,6 @@ const PERSISTENT_WORKSPACE_KEY_FUNCTIONS: Array<(workspaceId: string) => string> getReviewSearchStateKey, getReviewsKey, getAutoCompactionEnabledKey, - getStatusStateKey, // Note: getAutoCompactionThresholdKey is per-model, not per-workspace ]; diff --git a/src/common/orpc/schemas/stream.ts b/src/common/orpc/schemas/stream.ts index 996e5a302e..de096ee9ed 100644 --- a/src/common/orpc/schemas/stream.ts +++ b/src/common/orpc/schemas/stream.ts @@ -246,6 +246,15 @@ export const QueuedMessageChangedEventSchema = z.object({ imageParts: z.array(ImagePartSchema).optional(), }); +export const AgentStatusUpdateEventSchema = z.object({ + type: z.literal("agent-status-update"), + workspaceId: z.string(), + status: z.object({ + emoji: z.string().optional(), + message: z.string(), + url: z.string().url().optional(), + }), +}); export const RestoreToInputEventSchema = z.object({ type: z.literal("restore-to-input"), workspaceId: z.string(), @@ -277,6 +286,7 @@ export const WorkspaceChatMessageSchema = z.discriminatedUnion("type", [ // Usage and queue events UsageDeltaEventSchema, QueuedMessageChangedEventSchema, + AgentStatusUpdateEventSchema, RestoreToInputEventSchema, // Init events ...WorkspaceInitEventSchema.def.options, diff --git a/src/common/orpc/types.ts b/src/common/orpc/types.ts index 51a7759141..1e2c521ff4 100644 --- a/src/common/orpc/types.ts +++ b/src/common/orpc/types.ts @@ -115,6 +115,11 @@ export function isQueuedMessageChanged( return (msg as { type?: string }).type === "queued-message-changed"; } +export function isAgentStatusUpdate( + msg: WorkspaceChatMessage +): msg is Extract { + return (msg as { type?: string }).type === "agent-status-update"; +} export function isRestoreToInput( msg: WorkspaceChatMessage ): msg is Extract { diff --git a/src/common/types/tools.ts b/src/common/types/tools.ts index efabeb58ab..fcd9a500b7 100644 --- a/src/common/types/tools.ts +++ b/src/common/types/tools.ts @@ -221,9 +221,12 @@ export interface TodoWriteToolResult { // Status Set Tool Types export interface StatusSetToolArgs { - emoji: string; - message: string; - url?: string; + script: string; + /** + * Optional polling interval in seconds. + * If omitted, mux runs the script once. + */ + poll_interval_s?: number; } // Bash Output Tool Types (read incremental output from background processes) @@ -273,10 +276,8 @@ export type BashBackgroundListResult = export type StatusSetToolResult = | { + // status_set registers a poller and returns acknowledgement. success: true; - emoji: string; - message: string; - url?: string; } | { success: false; diff --git a/src/common/utils/status/parseAgentStatus.test.ts b/src/common/utils/status/parseAgentStatus.test.ts new file mode 100644 index 0000000000..8bbc293539 --- /dev/null +++ b/src/common/utils/status/parseAgentStatus.test.ts @@ -0,0 +1,33 @@ +import { describe, it, expect } from "bun:test"; +import { STATUS_MESSAGE_MAX_LENGTH } from "@/common/constants/toolLimits"; +import { parseAgentStatusFromLine } from "./parseAgentStatus"; + +describe("parseAgentStatusFromLine", () => { + it("extracts leading emoji and first URL, removing URL from message", () => { + const parsed = parseAgentStatusFromLine( + "πŸš€ PR #123 waiting for CI https://github.com/example/repo/pull/123", + STATUS_MESSAGE_MAX_LENGTH + ); + + expect(parsed).toEqual({ + emoji: "πŸš€", + message: "PR #123 waiting for CI", + url: "https://github.com/example/repo/pull/123", + }); + }); + + it("does not treat an emoji as leading emoji when not followed by whitespace", () => { + const parsed = parseAgentStatusFromLine("βœ…Done", STATUS_MESSAGE_MAX_LENGTH); + expect(parsed).toEqual({ message: "βœ…Done" }); + }); + + it("truncates after URL extraction", () => { + const long = `βœ… ${"a".repeat(STATUS_MESSAGE_MAX_LENGTH + 20)} https://example.com/pr/1`; + const parsed = parseAgentStatusFromLine(long, STATUS_MESSAGE_MAX_LENGTH); + + expect(parsed.emoji).toBe("βœ…"); + expect(parsed.url).toBe("https://example.com/pr/1"); + expect(parsed.message.length).toBe(STATUS_MESSAGE_MAX_LENGTH); + expect(parsed.message.endsWith("…")).toBe(true); + }); +}); diff --git a/src/common/utils/status/parseAgentStatus.ts b/src/common/utils/status/parseAgentStatus.ts new file mode 100644 index 0000000000..29df1bb093 --- /dev/null +++ b/src/common/utils/status/parseAgentStatus.ts @@ -0,0 +1,83 @@ +export interface ParsedAgentStatus { + emoji?: string; + message: string; + url?: string; +} + +// Basic URL matcher. Intentionally simple: we only need the first URL to make it clickable. +const URL_REGEX = /https?:\/\/[^\s"'<>]+/i; + +function truncateMessage(message: string, maxLength: number): string { + if (message.length <= maxLength) { + return message; + } + // Truncate to maxLength-1 and add ellipsis (total = maxLength) + return message.slice(0, maxLength - 1) + "…"; +} + +/** + * Returns true if `str` is a single emoji grapheme cluster. + * Uses Intl.Segmenter so variation selectors / skin tones count as a single cluster. + */ +function isSingleEmojiGrapheme(str: string): boolean { + if (!str) return false; + + const segmenter = new Intl.Segmenter("en", { granularity: "grapheme" }); + const segments = [...segmenter.segment(str)]; + if (segments.length !== 1) return false; + + const emojiRegex = /^[\p{Emoji_Presentation}\p{Extended_Pictographic}]/u; + return emojiRegex.test(segments[0].segment); +} + +function sanitizeUrl(url: string): string { + // Trim common trailing punctuation that often follows URLs in prose. + return url.replace(/[\])}.,;:!?]+$/g, ""); +} + +/** + * Parse a single-line status string into { emoji?, message, url? }. + * + * Parsing order matters: + * 1) Extract URL and remove it from the message (prevents redundancy in UI) + * 2) Extract leading emoji (if present) + * 3) Truncate AFTER extraction + */ +export function parseAgentStatusFromLine(rawLine: string, maxLength: number): ParsedAgentStatus { + let line = rawLine.trim(); + + // URL extraction (first URL only) + let url: string | undefined; + const urlMatch = URL_REGEX.exec(line); + if (urlMatch) { + url = sanitizeUrl(urlMatch[0]); + // Remove the matched substring (not the sanitized version) to keep surrounding punctuation intact. + const index = urlMatch.index; + line = (line.slice(0, index) + line.slice(index + urlMatch[0].length)) + .replace(/\s+/g, " ") + .trim(); + } + + // Leading emoji extraction + let emoji: string | undefined; + if (line.length > 0) { + const segmenter = new Intl.Segmenter("en", { granularity: "grapheme" }); + const first = segmenter.segment(line)[Symbol.iterator]().next(); + if (!first.done) { + const firstCluster = first.value.segment; + const rest = line.slice(first.value.index + firstCluster.length); + if (isSingleEmojiGrapheme(firstCluster) && /^\s/.test(rest)) { + emoji = firstCluster; + line = rest.trim(); + } + } + } + + const message = truncateMessage(line, maxLength); + + return { + ...(emoji ? { emoji } : {}), + message, + ...(url ? { url } : {}), + }; +} diff --git a/src/common/utils/tools/toolDefinitions.ts b/src/common/utils/tools/toolDefinitions.ts index 522eee72f5..5b6907a99e 100644 --- a/src/common/utils/tools/toolDefinitions.ts +++ b/src/common/utils/tools/toolDefinitions.ts @@ -10,7 +10,6 @@ import { BASH_HARD_MAX_LINES, BASH_MAX_LINE_BYTES, BASH_MAX_TOTAL_BYTES, - STATUS_MESSAGE_MAX_LENGTH, WEB_FETCH_MAX_OUTPUT_BYTES, } from "@/common/constants/toolLimits"; import { TOOL_EDIT_WARNING } from "@/common/types/tools"; @@ -203,36 +202,39 @@ export const TOOL_DEFINITIONS = { }, status_set: { description: - "Set a status indicator to show what Assistant is currently doing. The status is set IMMEDIATELY \n" + - "when this tool is called, even before other tool calls complete.\n" + + "Set a workspace status by registering a script that mux executes (optionally on an interval).\n" + + "The script runs in the workspace directory and should print a single line describing the status.\n" + "\n" + - "WHEN TO SET STATUS:\n" + - "- Set status when beginning concrete work (file edits, running tests, executing commands)\n" + - "- Update status as work progresses through distinct phases\n" + - "- Set a final status after completion, only claim success when certain (e.g., after confirming checks passed)\n" + - "- DO NOT set status during initial exploration, file reading, or planning phases\n" + + "OUTPUT FORMAT (stdout):\n" + + "- Optional leading emoji + whitespace, then message text\n" + + "- Include a URL anywhere in the line (e.g., a PR URL). mux will extract the first URL and show it as a clickable link.\n" + + " The URL is removed from the displayed message to avoid redundancy.\n" + "\n" + - "The status is cleared when a new user message comes in. Validate your approach is feasible \n" + - "before setting status - failed tool calls after setting status indicate premature commitment.\n" + + "POLLING:\n" + + "- poll_interval_s omitted: runs once (static status via `echo`)\n" + + "- poll_interval_s set: re-runs the script to keep the status fresh\n" + "\n" + - "URL PARAMETER:\n" + - "- Optional 'url' parameter links to external resources (e.g., PR URL: 'https://github.com/owner/repo/pull/123')\n" + - "- Prefer stable URLs that don't change often - saving the same URL twice is a no-op\n" + - "- URL persists until replaced by a new status with a different URL", + "BEST PRACTICES:\n" + + "- For long-running / multi-step work, prefer poll_interval_s (e.g., 5-15s) and a *dynamic* script so the user sees progress without extra chat messages.\n" + + "- Keep the script fast and single-line (it runs with a short timeout).\n" + + "- Include the most relevant URL (PR, CI run, etc.) so it's always clickable.\n" + + "\n" + + "NOTE: Workspace status persists after streaming completes and is not cleared on new user turns.", schema: z .object({ - emoji: z.string().describe("A single emoji character representing the current activity"), - message: z - .string() - .describe( - `A brief description of the current activity (auto-truncated to ${STATUS_MESSAGE_MAX_LENGTH} chars with ellipsis if needed)` - ), - url: z + script: z .string() - .url() + .min(1) + .describe("Shell script to execute. Print a single status line to stdout."), + poll_interval_s: z + .number() + .int() + .min(1) + .max(30) .optional() .describe( - "Optional URL to external resource with more details (e.g., Pull Request URL). The URL persists and is displayed to the user for easy access." + "Optional polling interval in seconds. If omitted, runs once. " + + "Use this for tasks that take time (builds/tests/CI wait) so the status updates automatically." ), }) .strict(), diff --git a/src/common/utils/tools/tools.ts b/src/common/utils/tools/tools.ts index 94c19a2ec0..e7bc4d4275 100644 --- a/src/common/utils/tools/tools.ts +++ b/src/common/utils/tools/tools.ts @@ -16,6 +16,7 @@ import { log } from "@/node/services/log"; import type { Runtime } from "@/node/runtime/Runtime"; import type { InitStateManager } from "@/node/services/initStateManager"; import type { BackgroundProcessManager } from "@/node/services/backgroundProcessManager"; +import type { StatusSetService } from "@/node/services/statusSetService"; import type { UIMode } from "@/common/types/mode"; import type { FileState } from "@/node/services/agentSession"; @@ -43,8 +44,19 @@ export interface ToolConfiguration { mode?: UIMode; /** Plan file path - only this file can be edited in plan mode */ planFilePath?: string; + /** + * Optional hook for tools to emit additional AIService events. + * Used for non-stream tool side channels like agent-status updates. + */ + emitAIEvent?: (event: string, payload: unknown) => void; /** Workspace ID for tracking background processes and plan storage */ workspaceId?: string; + /** + * StatusSetService - persists and rehydrates status_set polling across restarts. + * Provided by AIService so the status_set tool stays thin. + */ + statusSetService?: StatusSetService; + /** Callback to record file state for external edit detection (plan files) */ recordFileState?: (filePath: string, state: FileState) => void; } @@ -120,6 +132,7 @@ export async function getToolsForModel( bash: wrap(createBashTool(config)), bash_output: wrap(createBashOutputTool(config)), bash_background_list: wrap(createBashBackgroundListTool(config)), + status_set: wrap(createStatusSetTool(config)), bash_background_terminate: wrap(createBashBackgroundTerminateTool(config)), web_fetch: wrap(createWebFetchTool(config)), }; @@ -129,7 +142,6 @@ export async function getToolsForModel( propose_plan: createProposePlanTool(config), todo_write: createTodoWriteTool(config), todo_read: createTodoReadTool(config), - status_set: createStatusSetTool(config), }; // Base tools available for all models diff --git a/src/node/orpc/router.ts b/src/node/orpc/router.ts index 9a2cf18db3..7e6bd5c360 100644 --- a/src/node/orpc/router.ts +++ b/src/node/orpc/router.ts @@ -455,6 +455,14 @@ export const router = (authToken?: string) => { push(message); }); + // Rehydrate persisted status_set polling and send a snapshot for renderer reloads. + // (agent-status-update is not persisted in chat.jsonl) + await context.workspaceService.ensureStatusSetRunning(input.workspaceId); + const statusSnapshot = context.workspaceService.getStatusSetSnapshot(input.workspaceId); + if (statusSnapshot) { + push(statusSnapshot); + } + // 2. Replay history (sends caught-up at the end) await session.replayHistory(({ message }) => { push(message); diff --git a/src/node/services/agentSession.ts b/src/node/services/agentSession.ts index 57dffc0fbc..bfb6c60549 100644 --- a/src/node/services/agentSession.ts +++ b/src/node/services/agentSession.ts @@ -606,6 +606,7 @@ export class AgentSession { }); forward("reasoning-delta", (payload) => this.emitChatEvent(payload)); forward("reasoning-end", (payload) => this.emitChatEvent(payload)); + forward("agent-status-update", (payload) => this.emitChatEvent(payload)); forward("usage-delta", (payload) => this.emitChatEvent(payload)); forward("stream-abort", (payload) => this.emitChatEvent(payload)); diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 373346c32c..8a1986e4a9 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -26,6 +26,7 @@ import { createRuntime } from "@/node/runtime/runtimeFactory"; import { getMuxEnv, getRuntimeType } from "@/node/runtime/initHook"; import { secretsToRecord } from "@/common/types/secrets"; import type { MuxProviderOptions } from "@/common/types/providerOptions"; +import { StatusSetService } from "@/node/services/statusSetService"; import type { BackgroundProcessManager } from "@/node/services/backgroundProcessManager"; import type { FileState, EditedFileAttachment } from "@/node/services/agentSession"; import { log } from "./log"; @@ -278,6 +279,7 @@ export class AIService extends EventEmitter { private readonly initStateManager: InitStateManager; private readonly mockModeEnabled: boolean; private readonly mockScenarioPlayer?: MockScenarioPlayer; + private readonly statusSetService: StatusSetService; private readonly backgroundProcessManager?: BackgroundProcessManager; constructor( @@ -296,6 +298,9 @@ export class AIService extends EventEmitter { this.historyService = historyService; this.partialService = partialService; this.initStateManager = initStateManager; + this.statusSetService = new StatusSetService(this.config, (event, payload) => + this.emit(event, payload) + ); this.backgroundProcessManager = backgroundProcessManager; this.streamManager = new StreamManager(historyService, partialService, sessionUsageService); void this.ensureSessionsDir(); @@ -362,6 +367,17 @@ export class AIService extends EventEmitter { } } + async ensureStatusSetRunning(workspaceId: string): Promise { + await this.statusSetService.ensureRunning(workspaceId); + } + + getStatusSetSnapshot(workspaceId: string) { + return this.statusSetService.getSnapshot(workspaceId); + } + + stopStatusSet(workspaceId: string): void { + this.statusSetService.stop(workspaceId); + } isMockModeEnabled(): boolean { return this.mockModeEnabled; } @@ -1153,6 +1169,8 @@ export class AIService extends EventEmitter { metadata.name ), runtimeTempDir, + statusSetService: this.statusSetService, + emitAIEvent: (event, payload) => this.emit(event, payload), backgroundProcessManager: this.backgroundProcessManager, // Plan/exec mode configuration for plan file access. // - read: plan file is readable in all modes (useful context) diff --git a/src/node/services/statusScriptPoller.ts b/src/node/services/statusScriptPoller.ts new file mode 100644 index 0000000000..367d8c73f7 --- /dev/null +++ b/src/node/services/statusScriptPoller.ts @@ -0,0 +1,135 @@ +import type { Runtime } from "@/node/runtime/Runtime"; +import { execBuffered } from "@/node/utils/runtime/helpers"; +import { STATUS_MESSAGE_MAX_LENGTH } from "@/common/constants/toolLimits"; +import { log } from "@/node/services/log"; +import { + parseAgentStatusFromLine, + type ParsedAgentStatus, +} from "@/common/utils/status/parseAgentStatus"; + +export interface StatusScriptPollerConfig { + workspaceId: string; + runtime: Runtime; + cwd: string; + env?: Record; + script: string; + pollIntervalMs: number; +} + +export class StatusScriptPoller { + private timer: NodeJS.Timeout | null = null; + private running = false; + private generation = 0; + private lastUrl: string | undefined; + private lastEmitted: ParsedAgentStatus | undefined; + private lastScript: string | undefined; + + constructor(private readonly onStatus: (status: ParsedAgentStatus) => void | Promise) {} + + stop(): void { + if (this.timer) { + clearInterval(this.timer); + this.timer = null; + } + } + + /** + * Set/replace the poller configuration. + * Always runs the script immediately once. + */ + set(config: StatusScriptPollerConfig): void { + this.generation++; + const gen = this.generation; + + // If the caller changes the script, don't leak the previous script's URL. + if (this.lastScript !== config.script) { + this.lastScript = config.script; + this.lastUrl = undefined; + this.lastEmitted = undefined; + } + + this.stop(); + + const run = () => { + // Intentionally fire-and-forget; runOnce is internally serialized via this.running. + this.runOnce(config, gen).catch(() => { + // Ignore status script errors; keep last known status. + }); + }; + + // Run immediately (even if pollIntervalMs === 0) + run(); + + if (config.pollIntervalMs > 0) { + this.timer = setInterval(run, config.pollIntervalMs); + } + } + + private statusesEqual(a: ParsedAgentStatus | undefined, b: ParsedAgentStatus): boolean { + return a?.emoji === b.emoji && a?.message === b.message && a?.url === b.url; + } + + private async runOnce(config: StatusScriptPollerConfig, gen: number): Promise { + if (this.running) { + return; + } + + this.running = true; + try { + const result = await execBuffered(config.runtime, config.script, { + cwd: config.cwd, + env: config.env, + timeout: 5, + }); + + // If config changed while we were executing, drop the update. + if (gen !== this.generation) { + return; + } + + const stdout = result.stdout ?? ""; + const stderr = result.stderr ?? ""; + + // Prefer stdout; fall back to stderr (useful for quick debugging scripts). + const text = stdout.trim().length > 0 ? stdout : stderr; + const firstNonEmptyLine = text + .split(/\r?\n/g) + .map((l) => l.trim()) + .find((l) => l.length > 0); + + if (!firstNonEmptyLine) { + return; + } + + const parsed = parseAgentStatusFromLine(firstNonEmptyLine, STATUS_MESSAGE_MAX_LENGTH); + + // Preserve last URL if subsequent updates omit it. + const url = parsed.url ?? this.lastUrl; + if (url) { + this.lastUrl = url; + } + + const nextStatus: ParsedAgentStatus = { + ...parsed, + ...(url ? { url } : {}), + }; + + // Avoid re-emitting the same status every poll interval. + if (this.statusesEqual(this.lastEmitted, nextStatus)) { + return; + } + this.lastEmitted = nextStatus; + + await this.onStatus(nextStatus); + } catch (error) { + // Ignore status script errors; keep last known status. + // Only log in debug mode to avoid noisy console output for user scripts. + log.debug("status_set: status script failed", { + workspaceId: config.workspaceId, + error: error instanceof Error ? error.message : String(error), + }); + } finally { + this.running = false; + } + } +} diff --git a/src/node/services/statusSetService.ts b/src/node/services/statusSetService.ts new file mode 100644 index 0000000000..1ea7803099 --- /dev/null +++ b/src/node/services/statusSetService.ts @@ -0,0 +1,259 @@ +import type { Runtime } from "@/node/runtime/Runtime"; +import { createRuntime } from "@/node/runtime/runtimeFactory"; +import { getMuxEnv, getRuntimeType } from "@/node/runtime/initHook"; +import { log } from "@/node/services/log"; +import { StatusScriptPoller } from "@/node/services/statusScriptPoller"; +import { SessionFileManager } from "@/node/utils/sessionFile"; + +import type { Config } from "@/node/config"; +import type { StatusSetToolArgs } from "@/common/types/tools"; +import type { FrontendWorkspaceMetadata } from "@/common/types/workspace"; +import type { RuntimeConfig } from "@/common/types/runtime"; +import { secretsToRecord } from "@/common/types/secrets"; +import type { Result } from "@/common/types/result"; +import { Ok, Err } from "@/common/types/result"; +import type { ParsedAgentStatus } from "@/common/utils/status/parseAgentStatus"; + +export type AgentStatus = ParsedAgentStatus; + +export interface AgentStatusUpdateEvent { + type: "agent-status-update"; + workspaceId: string; + status: AgentStatus; +} + +interface PersistedStatusSetState { + version: 1; + script: string; + poll_interval_s?: number; + lastStatus?: AgentStatus; + updatedAt: string; +} + +interface StatusSetRuntimeContext { + runtime: Runtime; + cwd: string; + env: Record; +} + +interface PollerState { + poller: StatusScriptPoller; + lastStatus?: AgentStatus; + lastPersisted?: PersistedStatusSetState; + runtimeContext?: StatusSetRuntimeContext; +} + +export class StatusSetService { + private readonly file: SessionFileManager; + private readonly pollersByWorkspaceId = new Map(); + + constructor( + private readonly config: Config, + private readonly emitAIEvent: (event: string, payload: unknown) => void + ) { + // Must be initialized in the constructor so this.config is available. + this.file = new SessionFileManager(this.config, "status_set.json"); + } + + getSnapshot(workspaceId: string): AgentStatusUpdateEvent | null { + const state = this.pollersByWorkspaceId.get(workspaceId); + if (!state?.lastStatus) { + return null; + } + return { type: "agent-status-update", workspaceId, status: state.lastStatus }; + } + + stop(workspaceId: string): void { + const state = this.pollersByWorkspaceId.get(workspaceId); + if (!state) return; + state.poller.stop(); + this.pollersByWorkspaceId.delete(workspaceId); + } + + /** + * Ensure any persisted status_set config is loaded and polling is started. + * + * This is what makes status_set robust to: + * - backend restarts (config is reloaded from ~/.mux/sessions/{workspaceId}/status_set.json) + * - renderer reloads (callers can request getSnapshot() on every subscribe) + */ + async ensureRunning(workspaceId: string): Promise { + const state = this.getOrCreateState(workspaceId); + + if (state.lastPersisted && state.runtimeContext) { + // Already configured and running (or attempted). + return; + } + + const persisted = await this.file.read(workspaceId); + if (!persisted) { + return; + } + + state.lastPersisted = persisted; + state.lastStatus = persisted.lastStatus; + + const ctxResult = await this.getRuntimeContextForWorkspace(workspaceId); + if (!ctxResult.success) { + log.debug("status_set: failed to create runtime context for persisted poller", { + workspaceId, + error: ctxResult.error, + }); + return; + } + + state.runtimeContext = ctxResult.data; + + log.debug("status_set: rehydrating", { + workspaceId, + poll_interval_s: persisted.poll_interval_s ?? 0, + }); + + state.poller.set({ + workspaceId, + runtime: state.runtimeContext.runtime, + cwd: state.runtimeContext.cwd, + env: state.runtimeContext.env, + script: persisted.script, + pollIntervalMs: (persisted.poll_interval_s ?? 0) * 1000, + }); + } + + async setFromToolCall(args: { + workspaceId: string; + toolArgs: StatusSetToolArgs; + runtime: Runtime; + cwd: string; + env: Record; + }): Promise> { + const persisted: PersistedStatusSetState = { + version: 1, + script: args.toolArgs.script, + ...(args.toolArgs.poll_interval_s ? { poll_interval_s: args.toolArgs.poll_interval_s } : {}), + // New script config: clear lastStatus so we don't show stale info on restart. + lastStatus: undefined, + updatedAt: new Date().toISOString(), + }; + + // Persist first so restarts are robust even if the app crashes after updating memory. + const writeResult = await this.file.write(args.workspaceId, persisted); + if (!writeResult.success) { + return Err(writeResult.error); + } + + const state = this.getOrCreateState(args.workspaceId); + state.lastPersisted = persisted; + state.lastStatus = undefined; + state.runtimeContext = { + runtime: args.runtime, + cwd: args.cwd, + env: args.env, + }; + + log.debug("status_set: configured", { + workspaceId: args.workspaceId, + poll_interval_s: persisted.poll_interval_s ?? 0, + }); + + state.poller.set({ + workspaceId: args.workspaceId, + runtime: args.runtime, + cwd: args.cwd, + env: args.env, + script: args.toolArgs.script, + pollIntervalMs: (args.toolArgs.poll_interval_s ?? 0) * 1000, + }); + + return Ok(undefined); + } + + private getOrCreateState(workspaceId: string): PollerState { + const existing = this.pollersByWorkspaceId.get(workspaceId); + if (existing) return existing; + + const poller = new StatusScriptPoller(async (status) => { + await this.onStatus(workspaceId, status); + }); + + const state: PollerState = { poller }; + this.pollersByWorkspaceId.set(workspaceId, state); + return state; + } + + private async onStatus(workspaceId: string, status: AgentStatus): Promise { + const state = this.pollersByWorkspaceId.get(workspaceId); + if (!state) return; + + state.lastStatus = status; + + // Persist the last known status so it can be shown immediately after restart. + if (state.lastPersisted) { + const nextPersisted: PersistedStatusSetState = { + ...state.lastPersisted, + lastStatus: status, + updatedAt: new Date().toISOString(), + }; + const result = await this.file.write(workspaceId, nextPersisted); + if (result.success) { + state.lastPersisted = nextPersisted; + } + } + + log.debug("status_set: emit", { + workspaceId, + message: status.message, + url: status.url, + }); + + this.emitAIEvent("agent-status-update", { + type: "agent-status-update", + workspaceId, + status, + }); + } + + private async getRuntimeContextForWorkspace( + workspaceId: string + ): Promise> { + const metadata = await this.getWorkspaceMetadata(workspaceId); + if (!metadata.success) { + return Err(metadata.error); + } + + const runtimeConfig: RuntimeConfig = + metadata.data.runtimeConfig ?? ({ type: "local", srcBaseDir: this.config.srcDir } as const); + + let runtime: Runtime; + try { + runtime = createRuntime(runtimeConfig, { projectPath: metadata.data.projectPath }); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + return Err(`Failed to create runtime: ${message}`); + } + + const cwd = runtime.getWorkspacePath(metadata.data.projectPath, metadata.data.name); + + const projectSecrets = this.config.getProjectSecrets(metadata.data.projectPath); + const env = { + ...getMuxEnv( + metadata.data.projectPath, + getRuntimeType(metadata.data.runtimeConfig), + metadata.data.name + ), + ...secretsToRecord(projectSecrets), + }; + + return Ok({ runtime, cwd, env }); + } + + private async getWorkspaceMetadata( + workspaceId: string + ): Promise> { + const allMetadata = await this.config.getAllWorkspaceMetadata(); + const found = allMetadata.find((m) => m.id === workspaceId); + if (!found) { + return Err(`Workspace metadata not found for ${workspaceId}`); + } + return Ok(found); + } +} diff --git a/src/node/services/tools/status_set.test.ts b/src/node/services/tools/status_set.test.ts index 05c176940c..dc195032ce 100644 --- a/src/node/services/tools/status_set.test.ts +++ b/src/node/services/tools/status_set.test.ts @@ -1,246 +1,108 @@ import { describe, it, expect } from "bun:test"; -import { createStatusSetTool } from "./status_set"; -import type { ToolConfiguration } from "@/common/utils/tools/tools"; -import { createRuntime } from "@/node/runtime/runtimeFactory"; +import * as fs from "fs/promises"; +import * as path from "path"; import type { ToolCallOptions } from "ai"; -import { STATUS_MESSAGE_MAX_LENGTH } from "@/common/constants/toolLimits"; - -describe("status_set tool validation", () => { - const mockConfig: ToolConfiguration = { - cwd: "/test", - runtime: createRuntime({ type: "local", srcBaseDir: "/tmp" }), - runtimeTempDir: "/tmp", - workspaceId: "test-workspace", - }; - +import { createRuntime } from "@/node/runtime/runtimeFactory"; +import type { ToolConfiguration } from "@/common/utils/tools/tools"; +import { createStatusSetTool } from "./status_set"; +import { StatusSetService } from "@/node/services/statusSetService"; +import { Config } from "@/node/config"; +import { TestTempDir } from "./testHelpers"; + +async function waitFor(predicate: () => boolean, timeoutMs = 500): Promise { + const start = Date.now(); + for (;;) { + if (predicate()) return; + if (Date.now() - start > timeoutMs) { + throw new Error("Timed out waiting for condition"); + } + await new Promise((r) => setTimeout(r, 10)); + } +} + +describe("status_set tool", () => { const mockToolCallOptions: ToolCallOptions = { toolCallId: "test-call-id", messages: [], }; - describe("emoji validation", () => { - it("should accept single emoji characters", async () => { - const tool = createStatusSetTool(mockConfig); - - const emojis = ["πŸ”", "πŸ“", "βœ…", "πŸš€", "⏳"]; - for (const emoji of emojis) { - const result = (await tool.execute!({ emoji, message: "Test" }, mockToolCallOptions)) as { - success: boolean; - emoji: string; - message: string; - }; - expect(result).toEqual({ success: true, emoji, message: "Test" }); - } - }); - - it("should accept emojis with variation selectors", async () => { - const tool = createStatusSetTool(mockConfig); - - // Emojis with variation selectors (U+FE0F) - const emojis = ["✏️", "βœ…", "➑️", "β˜€οΈ"]; - for (const emoji of emojis) { - const result = (await tool.execute!({ emoji, message: "Test" }, mockToolCallOptions)) as { - success: boolean; - emoji: string; - message: string; - }; - expect(result).toEqual({ success: true, emoji, message: "Test" }); - } - }); + it("registers a script status, persists it, and rehydrates snapshot", async () => { + using tempRoot = new TestTempDir("status-set-root"); + using tempProject = new TestTempDir("status-set-project"); - it("should accept emojis with skin tone modifiers", async () => { - const tool = createStatusSetTool(mockConfig); - - const emojis = ["πŸ‘‹πŸ»", "πŸ‘‹πŸ½", "πŸ‘‹πŸΏ"]; - for (const emoji of emojis) { - const result = (await tool.execute!({ emoji, message: "Test" }, mockToolCallOptions)) as { - success: boolean; - emoji: string; - message: string; - }; - expect(result).toEqual({ success: true, emoji, message: "Test" }); - } - }); + const emitted: Array<{ event: string; payload: unknown }> = []; - it("should reject multiple emojis", async () => { - const tool = createStatusSetTool(mockConfig); + const config = new Config(tempRoot.path); + // Ensure sessions dir exists for SessionFileManager + await fs.mkdir(config.sessionsDir, { recursive: true }); - const result1 = (await tool.execute!( - { emoji: "πŸ”πŸ“", message: "Test" }, - mockToolCallOptions - )) as { success: boolean; error: string }; - expect(result1.success).toBe(false); - expect(result1.error).toBe("emoji must be a single emoji character"); - - const result2 = (await tool.execute!( - { emoji: "βœ…βœ…", message: "Test" }, - mockToolCallOptions - )) as { success: boolean; error: string }; - expect(result2.success).toBe(false); - expect(result2.error).toBe("emoji must be a single emoji character"); - }); - - it("should reject text (non-emoji)", async () => { - const tool = createStatusSetTool(mockConfig); - - const result1 = (await tool.execute!( - { emoji: "a", message: "Test" }, - mockToolCallOptions - )) as { - success: boolean; - error: string; - }; - expect(result1.success).toBe(false); - expect(result1.error).toBe("emoji must be a single emoji character"); - - const result2 = (await tool.execute!( - { emoji: "abc", message: "Test" }, - mockToolCallOptions - )) as { success: boolean; error: string }; - expect(result2.success).toBe(false); - expect(result2.error).toBe("emoji must be a single emoji character"); - - const result3 = (await tool.execute!( - { emoji: "!", message: "Test" }, - mockToolCallOptions - )) as { - success: boolean; - error: string; - }; - expect(result3.success).toBe(false); - expect(result3.error).toBe("emoji must be a single emoji character"); + const statusSetService = new StatusSetService(config, (event, payload) => { + emitted.push({ event, payload }); }); - it("should reject empty emoji", async () => { - const tool = createStatusSetTool(mockConfig); + const workspaceId = "test-workspace"; - const result = (await tool.execute!({ emoji: "", message: "Test" }, mockToolCallOptions)) as { - success: boolean; - error: string; - }; - expect(result.success).toBe(false); - expect(result.error).toBe("emoji must be a single emoji character"); - }); - - it("should reject emoji with text", async () => { - const tool = createStatusSetTool(mockConfig); + // Tool calls run in a real runtime context. Use LocalRuntime for deterministic cwd. + const mockConfig: ToolConfiguration = { + cwd: tempProject.path, + runtime: createRuntime({ type: "local" }, { projectPath: tempProject.path }), + runtimeTempDir: tempProject.path, + workspaceId, + statusSetService, + }; - const result1 = (await tool.execute!( - { emoji: "πŸ”a", message: "Test" }, - mockToolCallOptions - )) as { success: boolean; error: string }; - expect(result1.success).toBe(false); - expect(result1.error).toBe("emoji must be a single emoji character"); + const tool = createStatusSetTool(mockConfig); - const result2 = (await tool.execute!( - { emoji: "xπŸ”", message: "Test" }, + expect( + await tool.execute!( + { + script: "echo 'πŸš€ PR #123 waiting https://github.com/example/repo/pull/123'", + }, mockToolCallOptions - )) as { success: boolean; error: string }; - expect(result2.success).toBe(false); - expect(result2.error).toBe("emoji must be a single emoji character"); - }); - }); - - describe("message validation", () => { - it(`should accept messages up to ${STATUS_MESSAGE_MAX_LENGTH} characters`, async () => { - const tool = createStatusSetTool(mockConfig); - - const result1 = (await tool.execute!( - { emoji: "βœ…", message: "a".repeat(STATUS_MESSAGE_MAX_LENGTH) }, - mockToolCallOptions - )) as { success: boolean; message: string }; - expect(result1.success).toBe(true); - expect(result1.message).toBe("a".repeat(STATUS_MESSAGE_MAX_LENGTH)); - - const result2 = (await tool.execute!( - { emoji: "βœ…", message: "Analyzing code structure" }, - mockToolCallOptions - )) as { success: boolean }; - expect(result2.success).toBe(true); - }); - - it(`should truncate messages longer than ${STATUS_MESSAGE_MAX_LENGTH} characters with ellipsis`, async () => { - const tool = createStatusSetTool(mockConfig); - - // Test with MAX_LENGTH + 1 characters - const result1 = (await tool.execute!( - { emoji: "βœ…", message: "a".repeat(STATUS_MESSAGE_MAX_LENGTH + 1) }, - mockToolCallOptions - )) as { success: boolean; message: string }; - expect(result1.success).toBe(true); - expect(result1.message).toBe("a".repeat(STATUS_MESSAGE_MAX_LENGTH - 1) + "…"); - expect(result1.message.length).toBe(STATUS_MESSAGE_MAX_LENGTH); - - // Test with longer message - const longMessage = - "This is a very long message that exceeds the 60 character limit and should be truncated"; - const result2 = (await tool.execute!( - { emoji: "βœ…", message: longMessage }, - mockToolCallOptions - )) as { success: boolean; message: string }; - expect(result2.success).toBe(true); - expect(result2.message).toBe(longMessage.slice(0, STATUS_MESSAGE_MAX_LENGTH - 1) + "…"); - expect(result2.message.length).toBe(STATUS_MESSAGE_MAX_LENGTH); + ) + ).toEqual({ success: true }); + + await waitFor(() => emitted.some((e) => e.event === "agent-status-update")); + + const found = emitted.find((e) => e.event === "agent-status-update"); + expect(found).toBeDefined(); + if (!found) { + throw new Error("Missing agent-status-update event"); + } + + const payload = found.payload as { + type: string; + workspaceId: string; + status: { emoji?: string; message: string; url?: string }; + }; + + expect(payload.type).toBe("agent-status-update"); + expect(payload.workspaceId).toBe(workspaceId); + expect(payload.status).toEqual({ + emoji: "πŸš€", + message: "PR #123 waiting", + url: "https://github.com/example/repo/pull/123", }); - it("should accept empty message", async () => { - const tool = createStatusSetTool(mockConfig); - - const result = (await tool.execute!({ emoji: "βœ…", message: "" }, mockToolCallOptions)) as { - success: boolean; - }; - expect(result.success).toBe(true); - }); - }); - - describe("url parameter", () => { - it("should accept valid URLs", async () => { - const tool = createStatusSetTool(mockConfig); - - const validUrls = [ - "https://github.com/owner/repo/pull/123", - "http://example.com", - "https://example.com/path/to/resource?query=param", - ]; - - for (const url of validUrls) { - const result = (await tool.execute!( - { emoji: "πŸ”", message: "Test", url }, - mockToolCallOptions - )) as { - success: boolean; - url: string; - }; - expect(result.success).toBe(true); - expect(result.url).toBe(url); - } - }); - - it("should work without URL parameter", async () => { - const tool = createStatusSetTool(mockConfig); - - const result = (await tool.execute!( - { emoji: "βœ…", message: "Test" }, - mockToolCallOptions - )) as { - success: boolean; - url?: string; - }; - expect(result.success).toBe(true); - expect(result.url).toBeUndefined(); - }); - - it("should omit URL from result when undefined", async () => { - const tool = createStatusSetTool(mockConfig); - - const result = (await tool.execute!( - { emoji: "βœ…", message: "Test", url: undefined }, - mockToolCallOptions - )) as { - success: boolean; - }; - expect(result.success).toBe(true); - expect("url" in result).toBe(false); + // Persisted state should include the script, and lastStatus for restart robustness. + const persistedPath = path.join(config.getSessionDir(workspaceId), "status_set.json"); + const persistedRaw = await fs.readFile(persistedPath, "utf-8"); + const persisted = JSON.parse(persistedRaw) as { + script: string; + lastStatus?: { message: string; url?: string; emoji?: string }; + }; + + expect(persisted.script).toContain("PR #123"); + expect(persisted.lastStatus).toEqual(payload.status); + + // Simulate restart: new service instance should load snapshot from disk. + const statusSetService2 = new StatusSetService(config, () => undefined); + await statusSetService2.ensureRunning(workspaceId); + + expect(statusSetService2.getSnapshot(workspaceId)).toEqual({ + type: "agent-status-update", + workspaceId, + status: payload.status, }); }); }); diff --git a/src/node/services/tools/status_set.ts b/src/node/services/tools/status_set.ts index 0a8b138cc8..e1fd16be06 100644 --- a/src/node/services/tools/status_set.ts +++ b/src/node/services/tools/status_set.ts @@ -1,77 +1,49 @@ import { tool } from "ai"; -import type { ToolFactory } from "@/common/utils/tools/tools"; +import type { ToolFactory, ToolConfiguration } from "@/common/utils/tools/tools"; import { TOOL_DEFINITIONS } from "@/common/utils/tools/toolDefinitions"; -import { STATUS_MESSAGE_MAX_LENGTH } from "@/common/constants/toolLimits"; -import type { StatusSetToolResult } from "@/common/types/tools"; +import type { StatusSetToolArgs, StatusSetToolResult } from "@/common/types/tools"; /** - * Validates that a string is a single emoji character - * Uses Intl.Segmenter to count grapheme clusters (handles variation selectors, skin tones, etc.) - */ -function isValidEmoji(str: string): boolean { - if (!str) return false; - - // Use Intl.Segmenter to count grapheme clusters (what users perceive as single characters) - // This properly handles emojis with variation selectors (like ✏️), skin tones, flags, etc. - const segmenter = new Intl.Segmenter("en", { granularity: "grapheme" }); - const segments = [...segmenter.segment(str)]; - - // Must be exactly one grapheme cluster - if (segments.length !== 1) { - return false; - } - - // Check if it's an emoji using Unicode properties - const emojiRegex = /^[\p{Emoji_Presentation}\p{Extended_Pictographic}]/u; - return emojiRegex.test(segments[0].segment); -} - -/** - * Truncates a message to a maximum length, adding an ellipsis if truncated - */ -function truncateMessage(message: string, maxLength: number): string { - if (message.length <= maxLength) { - return message; - } - // Truncate to maxLength-1 and add ellipsis (total = maxLength) - return message.slice(0, maxLength - 1) + "…"; -} - -/** - * Status set tool factory for AI assistant - * Creates a tool that allows the AI to set status indicator showing current activity + * status_set tool (script-based) * - * The status is displayed IMMEDIATELY when this tool is called, even before other - * tool calls complete. This prevents agents from prematurely declaring success - * (e.g., "PR checks passed") when operations are still pending. Agents should only - * set success status after confirming the outcome of long-running operations. + * Registers a status script that mux executes (optionally repeatedly) to keep the workspace status fresh. * - * @param config Required configuration (not used for this tool, but required by interface) + * Implementation note: + * - The tool itself is intentionally thin. + * - Persistence + rehydration is owned by StatusSetService (so status survives restarts/reloads). */ -export const createStatusSetTool: ToolFactory = () => { +export const createStatusSetTool: ToolFactory = (config: ToolConfiguration) => { return tool({ description: TOOL_DEFINITIONS.status_set.description, inputSchema: TOOL_DEFINITIONS.status_set.schema, - execute: ({ emoji, message, url }): Promise => { - // Validate emoji - if (!isValidEmoji(emoji)) { - return Promise.resolve({ - success: false, - error: "emoji must be a single emoji character", - }); + execute: async (args: StatusSetToolArgs): Promise => { + const workspaceId = config.workspaceId; + if (!workspaceId) { + return { success: false, error: "status_set requires workspaceId" }; } - // Truncate message if necessary - const truncatedMessage = truncateMessage(message, STATUS_MESSAGE_MAX_LENGTH); + if (!config.statusSetService) { + return { success: false, error: "status_set requires statusSetService" }; + } - // Tool execution is a no-op on the backend - // The status is tracked by StreamingMessageAggregator and displayed in the frontend - return Promise.resolve({ - success: true, - emoji, - message: truncatedMessage, - ...(url && { url }), + const env = { + ...(config.muxEnv ?? {}), + ...(config.secrets ?? {}), + }; + + const result = await config.statusSetService.setFromToolCall({ + workspaceId, + toolArgs: args, + runtime: config.runtime, + cwd: config.cwd, + env, }); + + if (!result.success) { + return { success: false, error: result.error }; + } + + return { success: true }; }, }); }; diff --git a/src/node/services/workspaceService.ts b/src/node/services/workspaceService.ts index b87d0175c9..bd3db0a6a1 100644 --- a/src/node/services/workspaceService.ts +++ b/src/node/services/workspaceService.ts @@ -22,11 +22,12 @@ import { fileExists } from "@/node/utils/runtime/fileExists"; import { expandTilde, expandTildeForSSH } from "@/node/runtime/tildeExpansion"; import type { PostCompactionExclusions } from "@/common/types/attachment"; -import type { - SendMessageOptions, - DeleteMessage, - ImagePart, - WorkspaceChatMessage, +import { + isAgentStatusUpdate, + type SendMessageOptions, + type DeleteMessage, + type ImagePart, + type WorkspaceChatMessage, } from "@/common/orpc/types"; import type { workspace as workspaceSchemas } from "@/common/orpc/schemas/api"; import type { z } from "zod"; @@ -48,7 +49,7 @@ import { createBashTool } from "@/node/services/tools/bash"; import type { BashToolResult } from "@/common/types/tools"; import { secretsToRecord } from "@/common/types/secrets"; -import { movePlanFile } from "@/node/utils/runtime/helpers"; +import { execBuffered, movePlanFile } from "@/node/utils/runtime/helpers"; /** Maximum number of retry attempts when workspace name collides */ const MAX_WORKSPACE_NAME_COLLISION_RETRIES = 3; @@ -168,6 +169,17 @@ export class WorkspaceService extends EventEmitter { void this.updateStreamingStatus(data.workspaceId, false); } }); + + // Forward agent-status-update events (from poller) to the appropriate session + this.aiService.on("agent-status-update", (data: unknown) => { + if (isAgentStatusUpdate(data as WorkspaceChatMessage)) { + const event = data as WorkspaceChatMessage & { workspaceId: string }; + const session = this.sessions.get(event.workspaceId); + if (session) { + session.emitChatEvent(event); + } + } + }); } private emitWorkspaceActivity( @@ -276,6 +288,16 @@ export class WorkspaceService extends EventEmitter { return session; } + async ensureStatusSetRunning(workspaceId: string): Promise { + await this.aiService.ensureStatusSetRunning(workspaceId); + } + + getStatusSetSnapshot(workspaceId: string): WorkspaceChatMessage | null { + return ( + (this.aiService.getStatusSetSnapshot(workspaceId) as WorkspaceChatMessage | null) ?? null + ); + } + public disposeSession(workspaceId: string): void { const session = this.sessions.get(workspaceId); if (!session) { @@ -289,6 +311,8 @@ export class WorkspaceService extends EventEmitter { this.sessionSubscriptions.delete(workspaceId); } + // Stop any status_set pollers for this workspace to avoid leaks. + this.aiService.stopStatusSet(workspaceId); session.dispose(); this.sessions.delete(workspaceId); } @@ -1086,7 +1110,7 @@ export class WorkspaceService extends EventEmitter { try { // Use exec to delete files since runtime doesn't have a deleteFile method // Delete both paths in one command for efficiency - await runtime.exec(`rm -f ${quotedPlanPath} ${quotedLegacyPlanPath}`, { + await execBuffered(runtime, `rm -f ${quotedPlanPath} ${quotedLegacyPlanPath}`, { cwd: metadata.projectPath, timeout: 10, });