Skip to content

Commit 14d06d3

Browse files
nikita-ashihminnikita-ashihmin
committed
Expose goal state and controls
1 parent d5737cb commit 14d06d3

13 files changed

Lines changed: 169 additions & 8 deletions

‎src/AcpExtensions.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import type {
77
} from "@agentclientprotocol/sdk";
88

99
export const LEGACY_SET_SESSION_MODEL_METHOD = "session/set_model";
10+
export const GOAL_CONTROL_METHOD = "session/goal_control";
1011

1112
export type LegacySessionModel = {
1213
modelId: string;
@@ -42,11 +43,13 @@ export type ExtMethodRequest =
4243
AuthenticationStatusRequest
4344
| AuthenticationLogoutRequest
4445
| LegacySetSessionModelExtRequest
46+
| GoalControlExtRequest
4547

4648
export function isExtMethodRequest(request: { method: string, params: Record<string, unknown> }): request is ExtMethodRequest {
4749
return request.method === "authentication/status"
4850
|| request.method === "authentication/logout"
49-
|| request.method === LEGACY_SET_SESSION_MODEL_METHOD;
51+
|| request.method === LEGACY_SET_SESSION_MODEL_METHOD
52+
|| request.method === GOAL_CONTROL_METHOD;
5053
}
5154

5255
export type AuthenticationStatusRequest = { method: "authentication/status", params: {} }
@@ -60,6 +63,16 @@ export type LegacySetSessionModelExtRequest = {
6063
params: LegacySetSessionModelRequest;
6164
}
6265

66+
export type GoalControlRequest = {
67+
sessionId: SessionId;
68+
action: "pause" | "clear";
69+
}
70+
71+
export type GoalControlExtRequest = {
72+
method: typeof GOAL_CONTROL_METHOD;
73+
params: GoalControlRequest;
74+
}
75+
6376
export async function legacySetSessionModel(
6477
connection: Pick<ClientContext, "request">,
6578
params: LegacySetSessionModelRequest,

‎src/CodexAcpClient.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import type {
3333
SkillsListResponse,
3434
SandboxPolicy,
3535
Thread,
36+
ThreadGoal,
3637
ThreadGoalStatus,
3738
ThreadSourceKind,
3839
TurnCompletedNotification,
@@ -422,6 +423,11 @@ export class CodexAcpClient {
422423
await this.codexClient.runCompact({threadId: sessionId});
423424
}
424425

426+
async getGoal(sessionId: string): Promise<ThreadGoal | null> {
427+
const response = await this.codexClient.threadGoalGet({threadId: sessionId});
428+
return response?.goal ?? null;
429+
}
430+
425431
async setGoal(
426432
sessionId: string,
427433
objective: string,

‎src/CodexAcpServer.ts‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ import {
4848
type LegacySessionModelState,
4949
type LegacySetSessionModelRequest,
5050
type LegacySetSessionModelResponse,
51+
GOAL_CONTROL_METHOD,
5152
isExtMethodRequest,
5253
LEGACY_SET_SESSION_MODEL_METHOD,
5354
} from "./AcpExtensions";
@@ -86,6 +87,8 @@ export interface ThreadGoalSnapshot {
8687
objective: string;
8788
status: ThreadGoalStatus;
8889
tokenBudget: number | null;
90+
timeUsedSeconds: number;
91+
controlMethod: typeof GOAL_CONTROL_METHOD;
8992
}
9093

9194
export interface SessionState {
@@ -255,6 +258,18 @@ export class CodexAcpServer {
255258
}
256259
case LEGACY_SET_SESSION_MODEL_METHOD:
257260
return await this.unstable_setSessionModel(this.parseLegacySetSessionModelParams(methodRequest.params));
261+
case GOAL_CONTROL_METHOD: {
262+
const sessionState = this.sessions.get(methodRequest.params.sessionId);
263+
if (!sessionState) {
264+
throw RequestError.invalidParams(undefined, `Unknown session: ${methodRequest.params.sessionId}`);
265+
}
266+
if (methodRequest.params.action === "pause") {
267+
await this.runWithProcessCheck(() => this.codexAcpClient.setGoalStatus(sessionState.sessionId, "paused"));
268+
} else {
269+
await this.runWithProcessCheck(() => this.codexAcpClient.clearGoal(sessionState.sessionId));
270+
}
271+
return {};
272+
}
258273
}
259274
}
260275

@@ -443,6 +458,9 @@ export class CodexAcpServer {
443458
}
444459

445460
this.publishAvailableCommandsAsync(sessionState);
461+
if ("sessionId" in request) {
462+
this.publishCurrentGoalAsync(sessionState);
463+
}
446464
const sessionModelState: LegacySessionModelState = this.createModelState(models, currentModelId);
447465
const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState();
448466

@@ -879,6 +897,33 @@ export class CodexAcpServer {
879897
void this.availableCommands.publish(sessionState);
880898
}
881899

900+
private publishCurrentGoalAsync(sessionState: SessionState): void {
901+
void this.publishCurrentGoal(sessionState).catch((err) => {
902+
logger.error(`Failed to publish current goal for session ${sessionState.sessionId}`, err);
903+
});
904+
}
905+
906+
private async publishCurrentGoal(sessionState: SessionState): Promise<void> {
907+
const goal = await this.runWithProcessCheck(() => this.codexAcpClient.getGoal(sessionState.sessionId));
908+
const snapshot: ThreadGoalSnapshot | null = goal === null ? null : {
909+
objective: goal.objective.trim(),
910+
status: goal.status,
911+
tokenBudget: goal.tokenBudget,
912+
timeUsedSeconds: goal.timeUsedSeconds,
913+
controlMethod: GOAL_CONTROL_METHOD,
914+
};
915+
sessionState.currentGoal = snapshot;
916+
const session = new ACPSessionConnection(this.connection, sessionState.sessionId);
917+
await session.update({
918+
sessionUpdate: "session_info_update",
919+
_meta: {
920+
codex: {
921+
goal: snapshot,
922+
},
923+
},
924+
});
925+
}
926+
882927
private findCurrentModel(models: Model[], currentModelId: string): Model | undefined {
883928
const modelId = ModelId.fromString(currentModelId);
884929
return models.find(m => m.id === modelId.model);
@@ -992,6 +1037,7 @@ export class CodexAcpServer {
9921037
}
9931038

9941039
await this.availableCommands.publish(sessionState);
1040+
await this.publishCurrentGoal(sessionState);
9951041
const sessionModelState: LegacySessionModelState = this.createModelState(models, currentModelId);
9961042
const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState();
9971043

‎src/CodexAppServerClient.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ import type {
3737
ThreadGoalClearedNotification,
3838
ThreadGoalClearParams,
3939
ThreadGoalClearResponse,
40+
ThreadGoalGetParams,
41+
ThreadGoalGetResponse,
4042
ThreadGoalSetParams,
4143
ThreadGoalSetResponse,
4244
ThreadLoadedListParams,
@@ -545,6 +547,10 @@ export class CodexAppServerClient {
545547
return await this.sendRequest({ method: "thread/goal/set", params: params });
546548
}
547549

550+
async threadGoalGet(params: ThreadGoalGetParams): Promise<ThreadGoalGetResponse> {
551+
return await this.sendRequest({ method: "thread/goal/get", params: params });
552+
}
553+
548554
async threadGoalClear(params: ThreadGoalClearParams): Promise<ThreadGoalClearResponse> {
549555
return await this.sendRequest({ method: "thread/goal/clear", params: params });
550556
}

‎src/CodexEventHandler.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type {
44
ServerNotification
55
} from "./app-server";
66
import type {SessionState, ThreadGoalSnapshot} from "./CodexAcpServer";
7+
import {GOAL_CONTROL_METHOD} from "./AcpExtensions";
78
import {type PlanEntry, RequestError} from "@agentclientprotocol/sdk";
89
import {ACPSessionConnection, type AcpClientConnection, type UpdateSessionEvent} from "./ACPSessionConnection";
910
import type {
@@ -283,6 +284,8 @@ export class CodexEventHandler {
283284
objective: event.goal.objective.trim(),
284285
status: event.goal.status,
285286
tokenBudget: event.goal.tokenBudget,
287+
timeUsedSeconds: event.goal.timeUsedSeconds,
288+
controlMethod: GOAL_CONTROL_METHOD,
286289
};
287290
}
288291

‎src/__tests__/CodexACPAgent/CodexAcpClient.test.ts‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import {AgentMode} from "../../AgentMode";
1616
import type {Model, ReviewStartResponse, ThreadGoal, TurnCompletedNotification, TurnStartParams} from "../../app-server/v2";
1717
import type {RateLimitsMap} from "../../RateLimitsMap";
1818
import {ModelId} from "../../ModelId";
19+
import {GOAL_CONTROL_METHOD} from "../../AcpExtensions";
1920

2021
describe('ACP server test', { timeout: 40_000 }, () => {
2122

@@ -1667,9 +1668,11 @@ describe('ACP server test', { timeout: 40_000 }, () => {
16671668
_meta: {
16681669
codex: {
16691670
goal: {
1670-
objective: "Ship the migration and keep tests green",
1671-
status: "active",
1672-
tokenBudget: null,
1671+
objective: "Ship the migration and keep tests green",
1672+
status: "active",
1673+
tokenBudget: null,
1674+
timeUsedSeconds: 0,
1675+
controlMethod: "session/goal_control",
16731676
},
16741677
},
16751678
},
@@ -2388,6 +2391,26 @@ describe('ACP server test', { timeout: 40_000 }, () => {
23882391
await expect(promptPromise).resolves.toMatchObject({stopReason: "cancelled"});
23892392
});
23902393

2394+
it('controls an active goal through the out-of-band session extension', async () => {
2395+
const { mockFixture, sessionState } = setupPromptFixture();
2396+
// @ts-expect-error - registering local session state for the extension request path
2397+
mockFixture.getCodexAcpAgent().sessions.set("session-id", sessionState);
2398+
const setStatusSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "setGoalStatus").mockResolvedValue(undefined);
2399+
const clearGoalSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "clearGoal").mockResolvedValue(undefined);
2400+
2401+
await expect(mockFixture.getCodexAcpAgent().extMethod(GOAL_CONTROL_METHOD, {
2402+
sessionId: "session-id",
2403+
action: "pause",
2404+
})).resolves.toEqual({});
2405+
await expect(mockFixture.getCodexAcpAgent().extMethod(GOAL_CONTROL_METHOD, {
2406+
sessionId: "session-id",
2407+
action: "clear",
2408+
})).resolves.toEqual({});
2409+
2410+
expect(setStatusSpy).toHaveBeenCalledWith("session-id", "paused");
2411+
expect(clearGoalSpy).toHaveBeenCalledWith("session-id");
2412+
});
2413+
23912414
it('suppresses the first routed goal notification after cancellation marks the turn stale', async () => {
23922415
const { mockFixture } = setupPromptFixture();
23932416
const codexAppServerClient = mockFixture.getCodexAppServerClient();

‎src/__tests__/CodexACPAgent/data/load-session-history.json‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,28 @@
8484
}
8585
]
8686
}
87+
{
88+
"method": "sessionUpdate",
89+
"args": [
90+
{
91+
"sessionId": "session-1",
92+
"update": {
93+
"sessionUpdate": "session_info_update",
94+
"_meta": {
95+
"codex": {
96+
"goal": {
97+
"objective": "Keep the restored migration green",
98+
"status": "paused",
99+
"tokenBudget": null,
100+
"timeUsedSeconds": 46,
101+
"controlMethod": "session/goal_control"
102+
}
103+
}
104+
}
105+
}
106+
}
107+
]
108+
}
87109
{
88110
"method": "sessionUpdate",
89111
"args": [

‎src/__tests__/CodexACPAgent/data/load-session-response-item-history-fallback.json‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,22 @@
8484
}
8585
]
8686
}
87+
{
88+
"method": "sessionUpdate",
89+
"args": [
90+
{
91+
"sessionId": "session-legacy",
92+
"update": {
93+
"sessionUpdate": "session_info_update",
94+
"_meta": {
95+
"codex": {
96+
"goal": null
97+
}
98+
}
99+
}
100+
}
101+
]
102+
}
87103
{
88104
"method": "sessionUpdate",
89105
"args": [

‎src/__tests__/CodexACPAgent/data/thread-goal-updated-multiline.json‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@
1010
"goal": {
1111
"objective": "First task\nSecond task",
1212
"status": "budgetLimited",
13-
"tokenBudget": 1000
13+
"tokenBudget": 1000,
14+
"timeUsedSeconds": 30,
15+
"controlMethod": "session/goal_control"
1416
}
1517
}
1618
}

‎src/__tests__/CodexACPAgent/data/thread-goal-updated.json‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@
1010
"goal": {
1111
"objective": "Ship the goal update",
1212
"status": "active",
13-
"tokenBudget": null
13+
"tokenBudget": null,
14+
"timeUsedSeconds": 12,
15+
"controlMethod": "session/goal_control"
1416
}
1517
}
1618
}

0 commit comments

Comments
 (0)