tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

refactor: goose 2 ui used acp session id (#8985)

Browse Source
This commit is contained in:
Lifei Zhou
2026-05-05 12:40:34 +10:00
committed by GitHub
parent 2fe4c3d0bb
commit de471bc2b2
30 changed files with 319 additions and 735 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/goose2/src/shared/api/__tests__/acp.test.ts
+84 -14
View File
@@ -1,13 +1,18 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
const mockLoadSession = vi.fn();
const mockNewSession = vi.fn();
const mockSetProvider = vi.fn();
const mockSetModel = vi.fn();
vi.mock("../acpApi", () => ({
listProviders: vi.fn(),
prompt: vi.fn(),
setModel: vi.fn(),
setModel: (...args: unknown[]) => mockSetModel(...args),
setProvider: (...args: unknown[]) => mockSetProvider(...args),
listSessions: vi.fn(),
loadSession: (...args: unknown[]) => mockLoadSession(...args),
newSession: (...args: unknown[]) => mockNewSession(...args),
exportSession: vi.fn(),
importSession: vi.fn(),
forkSession: vi.fn(),
@@ -29,29 +34,94 @@ describe("acpLoadSession", () => {
vi.resetModules();
});
it("restores the prior session mapping when replay loading fails", async () => {
it("restores the prior prepared session registration when replay loading fails", async () => {
mockLoadSession.mockRejectedValueOnce(new Error("load failed"));
const sessionTracker = await import("../acpSessionTracker");
const sessionRegistry = await import("../acpSessionRegistry");
const { acpLoadSession } = await import("../acp");
sessionTracker.registerSession(
"local-session",
"goose-session-1",
sessionRegistry.registerPreparedSession(
"acp-session-1",
"goose",
"/tmp/original",
);
await expect(
acpLoadSession("local-session", "goose-session-2", "/tmp/replay"),
acpLoadSession("acp-session-1", "/tmp/replay"),
).rejects.toThrow("load failed");
expect(sessionTracker.getGooseSessionId("local-session")).toBe(
"goose-session-1",
);
expect(sessionTracker.getLocalSessionId("goose-session-1")).toBe(
"local-session",
);
expect(sessionTracker.getLocalSessionId("goose-session-2")).toBeNull();
expect(sessionRegistry.isSessionPrepared("acp-session-1")).toBe(true);
});
});
describe("acpCreateSession", () => {
beforeEach(() => {
vi.clearAllMocks();
vi.resetModules();
});
it("uses the ACP session id as the UI session id", async () => {
mockNewSession.mockResolvedValue({ sessionId: "acp-session-1" });
const sessionRegistry = await import("../acpSessionRegistry");
const { acpCreateSession } = await import("../acp");
await expect(
acpCreateSession("openai", "/tmp/project", {
projectId: "project-1",
personaId: "persona-1",
modelId: "gpt-4.1",
}),
).resolves.toEqual({ sessionId: "acp-session-1" });
expect(mockNewSession).toHaveBeenCalledWith(
"/tmp/project",
"openai",
"project-1",
"persona-1",
);
expect(mockLoadSession).not.toHaveBeenCalled();
expect(mockSetProvider).toHaveBeenCalledWith("acp-session-1", "openai");
expect(mockSetModel).toHaveBeenCalledWith("acp-session-1", "gpt-4.1");
expect(sessionRegistry.isSessionPrepared("acp-session-1")).toBe(true);
});
});
describe("acpPrepareSession", () => {
beforeEach(() => {
vi.clearAllMocks();
vi.resetModules();
});
it("loads the existing ACP session instead of creating a replacement", async () => {
mockLoadSession.mockResolvedValue(undefined);
const sessionRegistry = await import("../acpSessionRegistry");
const { acpPrepareSession } = await import("../acp");
await expect(
acpPrepareSession("acp-session-1", "openai", "/tmp/project"),
).resolves.toBeUndefined();
expect(mockLoadSession).toHaveBeenCalledWith(
"acp-session-1",
"/tmp/project",
);
expect(mockNewSession).not.toHaveBeenCalled();
expect(mockSetProvider).toHaveBeenCalledWith("acp-session-1", "openai");
expect(sessionRegistry.isSessionPrepared("acp-session-1")).toBe(true);
});
it("surfaces load failures instead of creating a new ACP session", async () => {
mockLoadSession.mockRejectedValueOnce(new Error("missing session"));
const { acpPrepareSession } = await import("../acp");
await expect(
acpPrepareSession("acp-session-1", "openai", "/tmp/project"),
).rejects.toThrow("missing session");
expect(mockNewSession).not.toHaveBeenCalled();
expect(mockSetProvider).not.toHaveBeenCalled();
});
});
ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts
+26 -38
View File
@@ -11,12 +11,11 @@ import {
handleSessionNotification,
setActiveMessageId,
} from "../acpNotificationHandler";
import { registerSession } from "../acpSessionTracker";
import { registerPreparedSession } from "../acpSessionRegistry";
function createMcpAppPayload(): McpAppPayload {
return {
sessionId: "local-session",
gooseSessionId: "goose-session",
sessionId: "acp-session",
toolCallId: "tool-1",
toolCallTitle: "mcp_app_bench__inspect_host_info",
source: "toolCallUpdateMeta",
@@ -34,8 +33,7 @@ function createMcpAppPayload(): McpAppPayload {
describe("acpNotificationHandler", () => {
beforeEach(() => {
clearMessageTracking();
clearReplayBuffer("local-session");
clearReplayBuffer("goose-session");
clearReplayBuffer("acp-session");
useChatStore.setState({
messagesBySession: {},
sessionStateById: {},
@@ -49,16 +47,11 @@ describe("acpNotificationHandler", () => {
});
it("keeps tool calls that arrive before the first text chunk on the pending assistant message", async () => {
registerSession(
"local-session",
"goose-session",
"goose",
"/Users/aharvard",
);
setActiveMessageId("goose-session", "assistant-1");
registerPreparedSession("acp-session", "goose", "/Users/aharvard");
setActiveMessageId("acp-session", "assistant-1");
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call",
toolCallId: "tool-1",
@@ -67,7 +60,7 @@ describe("acpNotificationHandler", () => {
} as never);
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call_update",
toolCallId: "tool-1",
@@ -94,7 +87,7 @@ describe("acpNotificationHandler", () => {
} as never);
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "agent_message_chunk",
content: {
@@ -106,14 +99,13 @@ describe("acpNotificationHandler", () => {
await waitFor(() => {
const message =
useChatStore.getState().messagesBySession["local-session"]?.[0];
useChatStore.getState().messagesBySession["acp-session"]?.[0];
expect(message?.content.some((block) => block.type === "mcpApp")).toBe(
true,
);
});
const [message] =
useChatStore.getState().messagesBySession["local-session"];
const [message] = useChatStore.getState().messagesBySession["acp-session"];
expect(message.id).toBe("assistant-1");
expect(message.content.map((block) => block.type)).toEqual([
"toolRequest",
@@ -146,17 +138,17 @@ describe("acpNotificationHandler", () => {
text: "The Host Info inspector is now open.",
});
expect(
useChatStore.getState().getSessionRuntime("local-session")
useChatStore.getState().getSessionRuntime("acp-session")
.streamingMessageId,
).toBe("assistant-1");
});
it("preserves ACP tool kind and locations on tool requests", async () => {
registerSession("local-session", "goose-session", "goose", "/Users/test");
setActiveMessageId("goose-session", "assistant-1");
registerPreparedSession("acp-session", "goose", "/Users/test");
setActiveMessageId("acp-session", "assistant-1");
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call",
toolCallId: "tool-1",
@@ -168,7 +160,7 @@ describe("acpNotificationHandler", () => {
} as never);
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call_update",
toolCallId: "tool-1",
@@ -177,8 +169,7 @@ describe("acpNotificationHandler", () => {
},
} as never);
const [message] =
useChatStore.getState().messagesBySession["local-session"];
const [message] = useChatStore.getState().messagesBySession["acp-session"];
expect(message.content[0]).toMatchObject({
type: "toolRequest",
id: "tool-1",
@@ -190,16 +181,15 @@ describe("acpNotificationHandler", () => {
});
it("preserves structured tool output when ACP provides rawOutput", async () => {
registerSession(
"local-session",
"goose-session",
registerPreparedSession(
"acp-session",
"goose",
"/Users/aharvard/.goose/artifacts",
);
setActiveMessageId("goose-session", "assistant-1");
setActiveMessageId("acp-session", "assistant-1");
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call",
toolCallId: "tool-1",
@@ -208,7 +198,7 @@ describe("acpNotificationHandler", () => {
} as never);
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call_update",
toolCallId: "tool-1",
@@ -229,8 +219,7 @@ describe("acpNotificationHandler", () => {
},
} as never);
const [message] =
useChatStore.getState().messagesBySession["local-session"];
const [message] = useChatStore.getState().messagesBySession["acp-session"];
expect(message.content[1]).toMatchObject({
type: "toolResponse",
id: "tool-1",
@@ -244,7 +233,7 @@ describe("acpNotificationHandler", () => {
});
it("replay keeps tool and MCP app content on an assistant message when tool events arrive before text", async () => {
const replaySessionId = "replay-goose-session";
const replaySessionId = "replay-acp-session";
useChatStore.setState({
loadingSessionIds: new Set<string>([replaySessionId]),
});
@@ -339,7 +328,6 @@ describe("acpNotificationHandler", () => {
payload: {
...createMcpAppPayload(),
sessionId: replaySessionId,
gooseSessionId: replaySessionId,
},
});
});
@@ -442,8 +430,8 @@ describe("acpNotificationHandler", () => {
});
});
it("replay preserves gooseSessionId in MCP app payloads before tracker registration", async () => {
const replaySessionId = "replay-goose-session-2";
it("replay attaches MCP app payloads to tool-only assistant messages", async () => {
const replaySessionId = "replay-acp-session-2";
const replayCreated = 1_700_000_240;
useChatStore.setState({
loadingSessionIds: new Set<string>([replaySessionId]),
@@ -496,7 +484,7 @@ describe("acpNotificationHandler", () => {
expect(mcpAppBlock).toMatchObject({
type: "mcpApp",
payload: expect.objectContaining({
gooseSessionId: replaySessionId,
sessionId: replaySessionId,
}),
});
});
ui/goose2/src/shared/api/__tests__/acpSessionInfoUpdate.test.ts
-2
View File
@@ -28,7 +28,6 @@ describe("ACP session info updates", () => {
it("applies generated session info updates to non-user-named sessions", async () => {
useChatSessionStore.getState().addSession({
id: "goose-session-title",
acpSessionId: "goose-session-title",
title: "New Chat",
createdAt: "2026-01-01T00:00:00.000Z",
updatedAt: "2026-01-01T00:00:00.000Z",
@@ -62,7 +61,6 @@ describe("ACP session info updates", () => {
it("ignores generated titles for user-named sessions", async () => {
useChatSessionStore.getState().addSession({
id: "goose-session-user-title",
acpSessionId: "goose-session-user-title",
title: "My Custom Title",
createdAt: "2026-01-01T00:00:00.000Z",
updatedAt: "2026-01-01T00:00:00.000Z",
ui/goose2/src/shared/api/__tests__/acpToolCallStatus.test.ts
+8 -10
View File
@@ -9,13 +9,13 @@ import {
handleSessionNotification,
setActiveMessageId,
} from "../acpNotificationHandler";
import { registerSession } from "../acpSessionTracker";
import { registerPreparedSession } from "../acpSessionRegistry";
describe("ACP tool call status handling", () => {
beforeEach(() => {
clearMessageTracking();
clearReplayBuffer("replay-failed-tool-session");
clearReplayBuffer("goose-session");
clearReplayBuffer("acp-session");
useChatStore.setState({
messagesBySession: {},
sessionStateById: {},
@@ -76,16 +76,15 @@ describe("ACP tool call status handling", () => {
});
it("marks failed live tool updates as errors", async () => {
registerSession(
"local-session",
"goose-session",
registerPreparedSession(
"acp-session",
"goose",
"/Users/aharvard/.goose/artifacts",
);
setActiveMessageId("goose-session", "assistant-1");
setActiveMessageId("acp-session", "assistant-1");
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call",
toolCallId: "tool-1",
@@ -94,7 +93,7 @@ describe("ACP tool call status handling", () => {
} as never);
await handleSessionNotification({
sessionId: "goose-session",
sessionId: "acp-session",
update: {
sessionUpdate: "tool_call_update",
toolCallId: "tool-1",
@@ -111,8 +110,7 @@ describe("ACP tool call status handling", () => {
},
} as never);
const [message] =
useChatStore.getState().messagesBySession["local-session"];
const [message] = useChatStore.getState().messagesBySession["acp-session"];
expect(message.content[0]).toMatchObject({
type: "toolRequest",
id: "tool-1",
ui/goose2/src/shared/api/acp.ts
+22 -47
View File
@@ -1,7 +1,7 @@
import type { ContentBlock } from "@agentclientprotocol/sdk";
import * as directAcp from "./acpApi";
import type { AcpSessionInfo } from "./acpApi";
import * as sessionTracker from "./acpSessionTracker";
import * as sessionRegistry from "./acpSessionRegistry";
import {
getCatalogEntry,
resolveAgentProviderCatalogId,
@@ -27,12 +27,9 @@ export interface AcpSendMessageOptions {
images?: [string, string][];
}
export interface AcpPrepareSessionOptions {
export interface AcpCreateSessionOptions {
personaId?: string;
projectId?: string;
}
export interface AcpCreateSessionOptions extends AcpPrepareSessionOptions {
modelId?: string | null;
}
@@ -85,8 +82,7 @@ export async function acpSendMessage(
const sid = sessionId.slice(0, 8);
const tStart = performance.now();
const gooseSessionId = sessionTracker.getGooseSessionId(sessionId, personaId);
if (!gooseSessionId) {
if (!sessionRegistry.isSessionPrepared(sessionId)) {
throw new Error("Session not prepared. Call acpPrepareSession first.");
}
@@ -113,7 +109,7 @@ export async function acpSendMessage(
}
const messageId = crypto.randomUUID();
setActiveMessageId(gooseSessionId, messageId);
setActiveMessageId(sessionId, messageId);
perfLog(
`[perf:send] ${sid} acpSendMessage → prompt(len=${prompt.length}, imgs=${images?.length ?? 0})`,
@@ -123,7 +119,7 @@ export async function acpSendMessage(
if (personaId) meta.personaId = personaId;
try {
await directAcp.prompt(
gooseSessionId,
sessionId,
content,
Object.keys(meta).length > 0 ? meta : undefined,
);
@@ -132,7 +128,7 @@ export async function acpSendMessage(
`[perf:send] ${sid} prompt() resolved in ${(tDone - tPrompt).toFixed(1)}ms (total acpSendMessage ${(tDone - tStart).toFixed(1)}ms)`,
);
} finally {
clearActiveMessageId(gooseSessionId);
clearActiveMessageId(sessionId);
}
}
@@ -141,24 +137,16 @@ export async function acpPrepareSession(
sessionId: string,
providerId: string,
workingDir: string,
options: AcpPrepareSessionOptions = {},
): Promise<string> {
): Promise<void> {
const sid = sessionId.slice(0, 8);
const t0 = performance.now();
perfLog(
`[perf:prepare] ${sid} acpPrepareSession start (provider=${providerId})`,
);
const gooseSessionId = await sessionTracker.prepareSession(
sessionId,
providerId,
workingDir,
options.personaId,
options.projectId,
);
await sessionRegistry.prepareSession(sessionId, providerId, workingDir);
perfLog(
`[perf:prepare] ${sid} acpPrepareSession done in ${(performance.now() - t0).toFixed(1)}ms`,
);
return gooseSessionId;
}
export async function acpCreateSession(
@@ -166,31 +154,26 @@ export async function acpCreateSession(
workingDir: string,
options: AcpCreateSessionOptions = {},
): Promise<{ sessionId: string }> {
const localSessionId = crypto.randomUUID();
const gooseSessionId = await acpPrepareSession(
localSessionId,
providerId,
const response = await directAcp.newSession(
workingDir,
options,
);
sessionTracker.registerSession(
gooseSessionId,
gooseSessionId,
providerId,
workingDir,
options.projectId,
options.personaId,
);
const sessionId = response.sessionId;
await directAcp.setProvider(sessionId, providerId);
sessionRegistry.registerPreparedSession(sessionId, providerId, workingDir);
if (options.modelId) {
await directAcp.setModel(gooseSessionId, options.modelId);
await directAcp.setModel(sessionId, options.modelId);
}
return { sessionId: gooseSessionId };
return { sessionId };
}
export async function acpSetModel(
sessionId: string,
modelId: string,
): Promise<void> {
const gooseSessionId = sessionTracker.getGooseSessionId(sessionId);
return directAcp.setModel(gooseSessionId ?? sessionId, modelId);
return directAcp.setModel(sessionId, modelId);
}
export type { AcpSessionInfo };
@@ -223,21 +206,19 @@ export async function acpSearchSessions(
*/
export async function acpLoadSession(
sessionId: string,
gooseSessionId: string,
workingDir?: string,
): Promise<void> {
const effectiveWorkingDir = workingDir ?? "~";
const sid = sessionId.slice(0, 8);
const t0 = performance.now();
const rollbackSessionRegistration = sessionTracker.registerSession(
const rollbackSessionRegistration = sessionRegistry.registerPreparedSession(
sessionId,
gooseSessionId,
"goose",
effectiveWorkingDir,
);
try {
perfLog(`[perf:load] ${sid} acpLoadSession → client.loadSession`);
await directAcp.loadSession(gooseSessionId, effectiveWorkingDir);
await directAcp.loadSession(sessionId, effectiveWorkingDir);
perfLog(
`[perf:load] ${sid} client.loadSession resolved in ${(performance.now() - t0).toFixed(1)}ms`,
);
@@ -261,17 +242,11 @@ export async function acpImportSession(json: string): Promise<AcpSessionInfo> {
export async function acpDuplicateSession(
sessionId: string,
): Promise<AcpSessionInfo> {
const gooseSessionId =
sessionTracker.getGooseSessionId(sessionId) ?? sessionId;
return directAcp.forkSession(gooseSessionId);
return directAcp.forkSession(sessionId);
}
/** Cancel an in-progress ACP session so the backend stops streaming. */
export async function acpCancelSession(
sessionId: string,
personaId?: string,
): Promise<boolean> {
const gooseSessionId = sessionTracker.getGooseSessionId(sessionId, personaId);
await directAcp.cancelSession(gooseSessionId ?? sessionId);
export async function acpCancelSession(sessionId: string): Promise<boolean> {
await directAcp.cancelSession(sessionId);
return true;
}
ui/goose2/src/shared/api/acpNotificationHandler.test.ts
+6 -43
View File
@@ -1,11 +1,7 @@
import { beforeEach, describe, expect, it } from "vitest";
import type { SessionNotification } from "@agentclientprotocol/sdk";
import { useChatStore } from "@/features/chat/stores/chatStore";
import {
clearReplayBuffer,
getAndDeleteReplayBuffer,
} from "@/features/chat/hooks/replayBuffer";
import { registerSession } from "./acpSessionTracker";
import { clearReplayBuffer } from "@/features/chat/hooks/replayBuffer";
import {
clearMessageTracking,
handleSessionNotification,
@@ -14,8 +10,8 @@ import {
describe("acpNotificationHandler", () => {
beforeEach(() => {
clearMessageTracking();
clearReplayBuffer("draft-session-1");
clearReplayBuffer("draft-session-2");
clearReplayBuffer("acp-session-1");
clearReplayBuffer("acp-session-2");
useChatStore.setState({
messagesBySession: {},
sessionStateById: {},
@@ -28,9 +24,9 @@ describe("acpNotificationHandler", () => {
});
});
it("buffers usage updates until the local session mapping is registered", async () => {
it("applies usage updates to the ACP session id", async () => {
const notification = {
sessionId: "goose-session-1",
sessionId: "acp-session-1",
update: {
sessionUpdate: "usage_update",
used: 512,
@@ -40,42 +36,9 @@ describe("acpNotificationHandler", () => {
await handleSessionNotification(notification);
expect(
useChatStore.getState().sessionStateById["draft-session-1"],
).toBeUndefined();
expect(
useChatStore.getState().sessionStateById["goose-session-1"],
).toBeUndefined();
registerSession("draft-session-1", "goose-session-1", "goose", "/tmp");
const runtime = useChatStore
.getState()
.getSessionRuntime("draft-session-1");
const runtime = useChatStore.getState().getSessionRuntime("acp-session-1");
expect(runtime.tokenState.accumulatedTotal).toBe(512);
expect(runtime.tokenState.contextLimit).toBe(8192);
expect(runtime.hasUsageSnapshot).toBe(true);
});
it("does not buffer non-usage updates before the local session mapping exists", async () => {
const notification = {
sessionId: "goose-session-2",
update: {
sessionUpdate: "agent_message_chunk",
messageId: "message-1",
content: {
type: "text",
text: "hello from replay",
},
},
} as SessionNotification;
await handleSessionNotification(notification);
registerSession("draft-session-2", "goose-session-2", "goose", "/tmp");
expect(getAndDeleteReplayBuffer("draft-session-2")).toBeUndefined();
expect(
useChatStore.getState().messagesBySession["draft-session-2"],
).toBeUndefined();
});
});
ui/goose2/src/shared/api/acpNotificationHandler.ts
+24 -77
View File
@@ -31,14 +31,10 @@ import {
} from "./acpReplayAssistant";
import { getReplayCreated, getReplayMessageId } from "./acpReplayMetadata";
import { handleSessionInfoUpdate } from "./acpSessionInfoUpdate";
import {
getLocalSessionId,
subscribeToSessionRegistration,
} from "./acpSessionTracker";
import { getToolCallIdentity } from "./acpToolCallIdentity";
import { perfLog } from "@/shared/lib/perfLog";
// Pre-set message ID for the next live stream per goose session
// Pre-set message ID for the next live stream per session.
const presetMessageIds = new Map<string, string>();
// Per-session perf counters for replay/live streaming.
@@ -54,10 +50,6 @@ interface LivePerf {
chunkCount: number;
}
const livePerf = new Map<string, LivePerf>();
const pendingUsageUpdates = new Map<
string,
{ accumulatedTotal: number; contextLimit: number }
>();
const toolCallStatusFromUpdate = (status: string): ToolCallStatus =>
status === "failed" ? "error" : "completed";
@@ -108,33 +100,20 @@ function toolCallUpdatePatch(
};
}
subscribeToSessionRegistration((localSessionId, gooseSessionId) => {
const pendingUsage = pendingUsageUpdates.get(gooseSessionId);
if (!pendingUsage) {
return;
}
useChatStore.getState().updateTokenState(localSessionId, pendingUsage);
pendingUsageUpdates.delete(gooseSessionId);
});
export function setActiveMessageId(
gooseSessionId: string,
messageId: string,
): void {
presetMessageIds.set(gooseSessionId, messageId);
livePerf.set(gooseSessionId, {
export function setActiveMessageId(sessionId: string, messageId: string): void {
presetMessageIds.set(sessionId, messageId);
livePerf.set(sessionId, {
sendStartedAt: performance.now(),
firstChunkAt: null,
chunkCount: 0,
});
}
export function clearActiveMessageId(gooseSessionId: string): void {
presetMessageIds.delete(gooseSessionId);
const perf = livePerf.get(gooseSessionId);
export function clearActiveMessageId(sessionId: string): void {
presetMessageIds.delete(sessionId);
const perf = livePerf.get(sessionId);
if (perf) {
const sid = gooseSessionId.slice(0, 8);
const sid = sessionId.slice(0, 8);
const total = performance.now() - perf.sendStartedAt;
const ttft =
perf.firstChunkAt !== null
@@ -143,16 +122,14 @@ export function clearActiveMessageId(gooseSessionId: string): void {
perfLog(
`[perf:stream] ${sid} stream ended — ttft=${ttft}ms total=${total.toFixed(1)}ms chunks=${perf.chunkCount}`,
);
livePerf.delete(gooseSessionId);
livePerf.delete(sessionId);
}
}
export async function handleSessionNotification(
notification: SessionNotification,
): Promise<void> {
const gooseSessionId = notification.sessionId;
const localSessionId = getLocalSessionId(gooseSessionId);
const sessionId = localSessionId ?? gooseSessionId;
const sessionId = notification.sessionId;
const { update } = notification;
const isReplay = useChatStore.getState().loadingSessionIds.has(sessionId);
@@ -167,20 +144,20 @@ export async function handleSessionNotification(
}
perf.lastAt = now;
perf.count += 1;
handleReplay(sessionId, gooseSessionId, localSessionId, update);
handleReplay(sessionId, update);
} else {
const perf = livePerf.get(gooseSessionId);
const perf = livePerf.get(sessionId);
if (perf && update.sessionUpdate === "agent_message_chunk") {
perf.chunkCount += 1;
if (perf.firstChunkAt === null) {
perf.firstChunkAt = performance.now();
const sid = gooseSessionId.slice(0, 8);
const sid = sessionId.slice(0, 8);
perfLog(
`[perf:stream] ${sid} first agent_message_chunk at ttft=${(perf.firstChunkAt - perf.sendStartedAt).toFixed(1)}ms`,
);
}
}
handleLive(sessionId, gooseSessionId, localSessionId, update);
handleLive(sessionId, update);
}
}
@@ -202,12 +179,7 @@ function getChunkMessageId(update: SessionUpdate): string | null {
: null;
}
function handleReplay(
sessionId: string,
gooseSessionId: string,
localSessionId: string | null,
update: SessionUpdate,
): void {
function handleReplay(sessionId: string, update: SessionUpdate): void {
switch (update.sessionUpdate) {
case "agent_message_chunk": {
const msg = ensureReplayAssistantMessage(
@@ -331,7 +303,6 @@ function handleReplay(
update,
true,
{
gooseSessionId,
replayMessageId,
},
);
@@ -344,7 +315,7 @@ function handleReplay(
case "session_info_update":
case "config_option_update":
case "usage_update":
handleShared(sessionId, gooseSessionId, localSessionId, update);
handleShared(sessionId, update);
break;
default:
@@ -352,19 +323,13 @@ function handleReplay(
}
}
function handleLive(
sessionId: string,
gooseSessionId: string,
localSessionId: string | null,
update: SessionUpdate,
): void {
function handleLive(sessionId: string, update: SessionUpdate): void {
const store = useChatStore.getState();
switch (update.sessionUpdate) {
case "agent_message_chunk": {
const messageId = ensureLiveAssistantMessage(
sessionId,
gooseSessionId,
getChunkMessageId(update) ?? undefined,
);
@@ -376,7 +341,7 @@ function handleLive(
}
case "tool_call": {
const messageId = ensureLiveAssistantMessage(sessionId, gooseSessionId);
const messageId = ensureLiveAssistantMessage(sessionId);
const identity = getToolCallIdentity(update);
const toolRequest: ToolRequestContent = {
@@ -395,7 +360,7 @@ function handleLive(
}
case "tool_call_update": {
const messageId = ensureLiveAssistantMessage(sessionId, gooseSessionId);
const messageId = ensureLiveAssistantMessage(sessionId);
const identity = getToolCallIdentity(update);
const patch = toolCallUpdatePatch(update);
@@ -460,9 +425,6 @@ function handleLive(
toolRequest?.name ?? update.title ?? "",
update,
false,
{
gooseSessionId,
},
);
}
}
@@ -472,7 +434,7 @@ function handleLive(
case "session_info_update":
case "config_option_update":
case "usage_update":
handleShared(sessionId, gooseSessionId, localSessionId, update);
handleShared(sessionId, update);
break;
default:
@@ -480,12 +442,7 @@ function handleLive(
}
}
function handleShared(
sessionId: string,
gooseSessionId: string,
localSessionId: string | null,
update: SessionUpdate,
): void {
function handleShared(sessionId: string, update: SessionUpdate): void {
switch (update.sessionUpdate) {
case "session_info_update": {
handleSessionInfoUpdate(sessionId, update);
@@ -535,15 +492,7 @@ function handleShared(
case "usage_update": {
const usage = update as SessionUpdate & { sessionUpdate: "usage_update" };
if (!localSessionId) {
pendingUsageUpdates.set(gooseSessionId, {
accumulatedTotal: usage.used,
contextLimit: usage.size,
});
break;
}
useChatStore.getState().updateTokenState(localSessionId, {
useChatStore.getState().updateTokenState(sessionId, {
accumulatedTotal: usage.used,
contextLimit: usage.size,
});
@@ -562,7 +511,6 @@ function findStreamingMessageId(sessionId: string): string | null {
function ensureLiveAssistantMessage(
sessionId: string,
gooseSessionId: string,
preferredMessageId?: string | null,
): string {
const store = useChatStore.getState();
@@ -578,7 +526,7 @@ function ensureLiveAssistantMessage(
const messageId =
preferredMessageId ??
presetMessageIds.get(gooseSessionId) ??
presetMessageIds.get(sessionId) ??
existingStreamingMessageId ??
crypto.randomUUID();
@@ -598,14 +546,13 @@ function ensureLiveAssistantMessage(
store.setPendingAssistantProvider(sessionId, null);
store.setStreamingMessageId(sessionId, messageId);
clearActiveMessageId(gooseSessionId);
clearActiveMessageId(sessionId);
return messageId;
}
export function clearMessageTracking(): void {
presetMessageIds.clear();
pendingUsageUpdates.clear();
clearReplayAssistantTracking();
}
ui/goose2/src/shared/api/acpSessionRegistry.ts
+80
View File
@@ -0,0 +1,80 @@
import * as acpApi from "./acpApi";
import { perfLog } from "@/shared/lib/perfLog";
interface PreparedSession {
providerId: string;
workingDir: string;
}
const prepared = new Map<string, PreparedSession>();
export async function prepareSession(
sessionId: string,
providerId: string,
workingDir: string,
): Promise<void> {
const sid = sessionId.slice(0, 8);
const existing = prepared.get(sessionId);
if (existing) {
const tReuse = performance.now();
let changed = false;
if (existing.workingDir !== workingDir) {
await acpApi.updateWorkingDir(sessionId, workingDir);
existing.workingDir = workingDir;
changed = true;
}
if (existing.providerId !== providerId) {
const tProv = performance.now();
await acpApi.setProvider(sessionId, providerId);
perfLog(
`[perf:prepare] ${sid} reuse setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms`,
);
existing.providerId = providerId;
changed = true;
}
perfLog(
`[perf:prepare] ${sid} reuse existing session (updates=${changed}) in ${(performance.now() - tReuse).toFixed(1)}ms`,
);
return;
}
const tLoad = performance.now();
await acpApi.loadSession(sessionId, workingDir);
perfLog(
`[perf:prepare] ${sid} registry loadSession ok in ${(performance.now() - tLoad).toFixed(1)}ms`,
);
const tProv = performance.now();
await acpApi.setProvider(sessionId, providerId);
perfLog(
`[perf:prepare] ${sid} registry setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms`,
);
const entry = { providerId, workingDir };
prepared.set(sessionId, entry);
return;
}
export function isSessionPrepared(sessionId: string): boolean {
return prepared.has(sessionId);
}
export function registerPreparedSession(
sessionId: string,
providerId: string,
workingDir: string,
): () => void {
const previousEntry = prepared.get(sessionId);
const entry = { providerId, workingDir };
prepared.set(sessionId, entry);
return () => {
prepared.delete(sessionId);
if (previousEntry) {
prepared.set(sessionId, previousEntry);
}
};
}
ui/goose2/src/shared/api/acpSessionTracker.ts
-206
View File
@@ -1,206 +0,0 @@
import * as acpApi from "./acpApi";
import { perfLog } from "@/shared/lib/perfLog";
interface PreparedSession {
gooseSessionId: string;
providerId: string;
workingDir: string;
}
type SessionRegistrationListener = (
localSessionId: string,
gooseSessionId: string,
) => void;
const prepared = new Map<string, PreparedSession>();
const gooseToLocal = new Map<string, string>();
const registrationListeners = new Set<SessionRegistrationListener>();
function restoreGooseRegistration(
gooseSessionId: string,
localSessionId: string | undefined,
): void {
if (localSessionId === undefined) {
gooseToLocal.delete(gooseSessionId);
return;
}
gooseToLocal.set(gooseSessionId, localSessionId);
}
function makeKey(sessionId: string, personaId?: string): string {
if (personaId && personaId.length > 0) {
return `${sessionId}__${personaId}`;
}
return sessionId;
}
function notifySessionRegistered(
localSessionId: string,
gooseSessionId: string,
): void {
for (const listener of registrationListeners) {
listener(localSessionId, gooseSessionId);
}
}
export function subscribeToSessionRegistration(
listener: SessionRegistrationListener,
): () => void {
registrationListeners.add(listener);
return () => registrationListeners.delete(listener);
}
export async function prepareSession(
sessionId: string,
providerId: string,
workingDir: string,
personaId?: string,
projectId?: string,
): Promise<string> {
const sid = sessionId.slice(0, 8);
const key = makeKey(sessionId, personaId);
const existing = prepared.get(key) ?? prepared.get(sessionId);
if (existing) {
const tReuse = performance.now();
let changed = false;
if (existing.workingDir !== workingDir) {
await acpApi.updateWorkingDir(existing.gooseSessionId, workingDir);
existing.workingDir = workingDir;
changed = true;
}
if (existing.providerId !== providerId) {
const tProv = performance.now();
await acpApi.setProvider(existing.gooseSessionId, providerId);
perfLog(
`[perf:prepare] ${sid} reuse setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms (goose_sid=${existing.gooseSessionId.slice(0, 8)})`,
);
existing.providerId = providerId;
changed = true;
}
perfLog(
`[perf:prepare] ${sid} reuse existing session (updates=${changed}) in ${(performance.now() - tReuse).toFixed(1)}ms`,
);
return existing.gooseSessionId;
}
let gooseSessionId: string | null = null;
const tLoad = performance.now();
try {
await acpApi.loadSession(sessionId, workingDir);
gooseSessionId = sessionId;
perfLog(
`[perf:prepare] ${sid} tracker loadSession ok in ${(performance.now() - tLoad).toFixed(1)}ms`,
);
} catch {
perfLog(
`[perf:prepare] ${sid} tracker loadSession failed in ${(performance.now() - tLoad).toFixed(1)}ms → newSession`,
);
}
if (!gooseSessionId) {
const tNew = performance.now();
const response = await acpApi.newSession(
workingDir,
providerId,
projectId,
personaId,
);
gooseSessionId = response.sessionId;
perfLog(
`[perf:prepare] ${sid} tracker newSession done in ${(performance.now() - tNew).toFixed(1)}ms (goose_sid=${gooseSessionId.slice(0, 8)})`,
);
}
const gooseSid = gooseSessionId.slice(0, 8);
const tProv = performance.now();
await acpApi.setProvider(gooseSessionId, providerId);
perfLog(
`[perf:prepare] ${sid} tracker setProvider(${providerId}) in ${(performance.now() - tProv).toFixed(1)}ms (goose_sid=${gooseSid})`,
);
const entry = { gooseSessionId, providerId, workingDir };
prepared.set(key, entry);
prepared.set(sessionId, entry);
prepared.set(gooseSessionId, entry);
gooseToLocal.set(gooseSessionId, sessionId);
notifySessionRegistered(sessionId, gooseSessionId);
return gooseSessionId;
}
export function getGooseSessionId(
sessionId: string,
personaId?: string,
): string | null {
const key = makeKey(sessionId, personaId);
return (
prepared.get(key)?.gooseSessionId ??
prepared.get(sessionId)?.gooseSessionId ??
null
);
}
export function getLocalSessionId(gooseSessionId: string): string | null {
return gooseToLocal.get(gooseSessionId) ?? null;
}
export function registerSession(
sessionId: string,
gooseSessionId: string,
providerId: string,
workingDir: string,
): () => void {
const previousEntry = prepared.get(sessionId);
const previousGooseSessionLocal = gooseToLocal.get(gooseSessionId);
const previousSessionGooseLocal = previousEntry
? gooseToLocal.get(previousEntry.gooseSessionId)
: undefined;
const entry = { gooseSessionId, providerId, workingDir };
if (
previousEntry &&
previousEntry.gooseSessionId !== gooseSessionId &&
gooseToLocal.get(previousEntry.gooseSessionId) === sessionId
) {
gooseToLocal.delete(previousEntry.gooseSessionId);
}
prepared.set(sessionId, entry);
prepared.set(gooseSessionId, entry);
gooseToLocal.set(gooseSessionId, sessionId);
notifySessionRegistered(sessionId, gooseSessionId);
return () => {
prepared.delete(sessionId);
if (previousEntry) {
prepared.set(sessionId, previousEntry);
}
restoreGooseRegistration(gooseSessionId, previousGooseSessionLocal);
if (previousEntry && previousEntry.gooseSessionId !== gooseSessionId) {
restoreGooseRegistration(
previousEntry.gooseSessionId,
previousSessionGooseLocal,
);
}
};
}
export function unregisterSession(
sessionId: string,
gooseSessionId?: string,
): void {
const entry = prepared.get(sessionId);
prepared.delete(sessionId);
const resolvedGooseSessionId = gooseSessionId ?? entry?.gooseSessionId;
if (
resolvedGooseSessionId &&
gooseToLocal.get(resolvedGooseSessionId) === sessionId
) {
gooseToLocal.delete(resolvedGooseSessionId);
}
}
ui/goose2/src/shared/api/acpToolCallContent.ts
-2
View File
@@ -66,7 +66,6 @@ export function attachMcpAppPayload(
update: SessionUpdate,
isReplay: boolean,
options?: {
gooseSessionId?: string | null;
replayMessageId?: string | null;
},
): void {
@@ -75,7 +74,6 @@ export function attachMcpAppPayload(
toolCallId,
toolCallTitle,
update,
options?.gooseSessionId,
);
if (!payload) {
return;
ui/goose2/src/shared/api/mcpAppToolUpdate.ts
-4
View File
@@ -6,7 +6,6 @@ import type {
} from "@aaif/goose-sdk";
import type { SessionUpdate } from "@agentclientprotocol/sdk";
import type { McpAppPayload } from "@/shared/types/messages";
import { getGooseSessionId } from "./acpSessionTracker";
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null;
@@ -29,7 +28,6 @@ export function buildMcpAppPayloadFromToolUpdate(
toolCallId: string,
toolCallTitle: string,
update: SessionUpdate,
gooseSessionIdOverride?: string | null,
): McpAppPayload | null {
const payload = extractMcpAppPayload(update);
if (!payload) {
@@ -38,8 +36,6 @@ export function buildMcpAppPayloadFromToolUpdate(
return {
sessionId,
gooseSessionId:
gooseSessionIdOverride ?? getGooseSessionId(sessionId) ?? null,
toolCallId,
toolCallTitle,
source: "toolCallUpdateMeta",
ui/goose2/src/shared/types/messages.ts
-1
View File
@@ -118,7 +118,6 @@ export interface ToolResponseContent {
export interface McpAppPayload {
sessionId: string;
gooseSessionId: string | null;
toolCallId: string;
toolCallTitle: string;
source: "toolCallUpdateMeta";