Skip to content

Commit 6e5260c

Browse files
joshsnyjoshsny
committed
init resume and tree tracking
1 parent ab085b7 commit 6e5260c

9 files changed

Lines changed: 1484 additions & 26 deletions

File tree

‎packages/agent/src/acp-extensions.ts‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
/**
2+
* PostHog-specific ACP extensions.
3+
*
4+
* These follow the ACP extensibility model:
5+
* - Custom notification methods are prefixed with `_posthog/`
6+
* - Custom data can be attached via `_meta` fields
7+
*
8+
* Note: When using `extNotification()` from the ACP SDK, it automatically
9+
* adds an extra underscore prefix (e.g., `_posthog/tree_snapshot` becomes
10+
* `__posthog/tree_snapshot` in the log). Code that reads logs should handle both.
11+
*
12+
* See: https://agentclientprotocol.com/docs/extensibility
13+
*/
14+
15+
/**
16+
* Custom notification methods for PostHog-specific events.
17+
* Used with AgentSideConnection.extNotification() or Client.extNotification()
18+
*/
19+
export const POSTHOG_NOTIFICATIONS = {
20+
/** Git branch was created for a task */
21+
BRANCH_CREATED: "_posthog/branch_created",
22+
23+
/** Task run has started execution */
24+
RUN_STARTED: "_posthog/run_started",
25+
26+
/** Task has completed (success or failure) */
27+
TASK_COMPLETE: "_posthog/task_complete",
28+
29+
/** Error occurred during task execution */
30+
ERROR: "_posthog/error",
31+
32+
/** Console/log output from the agent */
33+
CONSOLE: "_posthog/console",
34+
35+
/** Maps a session ID to the underlying SDK session ID (for resumption) */
36+
SDK_SESSION: "_posthog/sdk_session",
37+
38+
/** Tree state snapshot captured (git tree hash + file archive) */
39+
TREE_SNAPSHOT: "_posthog/tree_snapshot",
40+
41+
/** Agent mode changed (interactive/background) */
42+
MODE_CHANGE: "_posthog/mode_change",
43+
44+
/** Request to resume a session from previous state */
45+
SESSION_RESUME: "_posthog/session/resume",
46+
47+
/** User message sent from client to agent */
48+
USER_MESSAGE: "_posthog/user_message",
49+
50+
/** Request to cancel current operation */
51+
CANCEL: "_posthog/cancel",
52+
53+
/** Request to close the session */
54+
CLOSE: "_posthog/close",
55+
56+
/** Agent status update (thinking, working, etc.) */
57+
STATUS: "_posthog/status",
58+
59+
/** Task-level notification (progress, milestones) */
60+
TASK_NOTIFICATION: "_posthog/task_notification",
61+
62+
/** Marks a boundary for log compaction */
63+
COMPACT_BOUNDARY: "_posthog/compact_boundary",
64+
} as const;
65+
66+
export type PostHogNotificationType =
67+
(typeof POSTHOG_NOTIFICATIONS)[keyof typeof POSTHOG_NOTIFICATIONS];
68+
69+
// --- Payload types for each notification ---
70+
71+
export interface BranchCreatedPayload {
72+
branch: string;
73+
}
74+
75+
export interface RunStartedPayload {
76+
sessionId: string;
77+
runId: string;
78+
taskId?: string;
79+
}
80+
81+
export interface TaskCompletePayload {
82+
sessionId: string;
83+
taskId: string;
84+
}
85+
86+
export interface ErrorNotificationPayload {
87+
sessionId: string;
88+
message: string;
89+
error?: unknown;
90+
}
91+
92+
export interface ConsoleNotificationPayload {
93+
sessionId: string;
94+
level: "debug" | "info" | "warn" | "error";
95+
message: string;
96+
}
97+
98+
export interface SdkSessionPayload {
99+
sessionId: string;
100+
sdkSessionId: string;
101+
}
102+
103+
export interface TreeSnapshotPayload {
104+
treeHash: string;
105+
baseCommit: string | null;
106+
archiveUrl?: string;
107+
filesChanged: string[];
108+
filesDeleted?: string[];
109+
timestamp: string;
110+
interrupted?: boolean;
111+
device?: {
112+
id: string;
113+
type: "local" | "cloud";
114+
name?: string;
115+
};
116+
}
117+
118+
export interface ModeChangePayload {
119+
mode: "interactive" | "background";
120+
previous_mode: "interactive" | "background";
121+
}
122+
123+
export interface SessionResumePayload {
124+
sessionId: string;
125+
fromSnapshot?: string;
126+
}
127+
128+
export interface UserMessagePayload {
129+
content: string;
130+
}
131+
132+
export interface StatusPayload {
133+
sessionId: string;
134+
status: string;
135+
message?: string;
136+
}
137+
138+
export interface TaskNotificationPayload {
139+
sessionId: string;
140+
type: string;
141+
message?: string;
142+
data?: Record<string, unknown>;
143+
}
144+
145+
export interface CompactBoundaryPayload {
146+
sessionId: string;
147+
timestamp: string;
148+
}
149+
150+
export type PostHogNotificationPayload =
151+
| BranchCreatedPayload
152+
| RunStartedPayload
153+
| TaskCompletePayload
154+
| ErrorNotificationPayload
155+
| ConsoleNotificationPayload
156+
| SdkSessionPayload
157+
| TreeSnapshotPayload
158+
| ModeChangePayload
159+
| SessionResumePayload
160+
| UserMessagePayload
161+
| StatusPayload
162+
| TaskNotificationPayload
163+
| CompactBoundaryPayload;

‎packages/agent/src/agent.ts‎

Lines changed: 145 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,31 @@
1+
import { POSTHOG_NOTIFICATIONS } from "./acp-extensions.js";
12
import {
23
createAcpConnection,
34
type InProcessAcpConnection,
45
} from "./adapters/acp-connection.js";
56
import { PostHogAPIClient } from "./posthog-api.js";
67
import { SessionLogWriter } from "./session-log-writer.js";
7-
import type { AgentConfig, TaskExecutionOptions } from "./types.js";
8+
import { TreeTracker } from "./tree-tracker.js";
9+
import type {
10+
AgentConfig,
11+
AgentMode,
12+
DeviceInfo,
13+
TaskExecutionOptions,
14+
TreeSnapshotEvent,
15+
} from "./types.js";
816
import { Logger } from "./utils/logger.js";
917

1018
export class Agent {
1119
private posthogAPI?: PostHogAPIClient;
1220
private logger: Logger;
1321
private acpConnection?: InProcessAcpConnection;
14-
private taskRunId?: string;
1522
private sessionLogWriter?: SessionLogWriter;
23+
private currentRunId?: string;
24+
private currentTaskId?: string;
25+
private treeTracker?: TreeTracker;
26+
private deviceInfo?: DeviceInfo;
27+
private agentMode: AgentMode = "interactive";
28+
private isExecutingTool: boolean = false;
1629
public debug: boolean;
1730

1831
constructor(config: AgentConfig) {
@@ -55,7 +68,29 @@ export class Agent {
5568
): Promise<InProcessAcpConnection> {
5669
await this._configureLlmGateway();
5770

58-
this.taskRunId = taskRunId;
71+
const cwd = options.repositoryPath;
72+
73+
// Use taskRunId as sessionId - they are the same identifier
74+
this.currentRunId = taskRunId;
75+
this.currentTaskId = taskId;
76+
77+
// Initialize TreeTracker for state capture if we have a repository path
78+
if (cwd) {
79+
this.treeTracker = new TreeTracker({
80+
repositoryPath: cwd,
81+
taskId,
82+
runId: taskRunId,
83+
apiClient: this.posthogAPI,
84+
logger: this.logger.child("TreeTracker"),
85+
});
86+
}
87+
88+
// Set device info for local mode
89+
this.deviceInfo = {
90+
id: `local-${process.pid}`,
91+
type: "local",
92+
name: process.env.HOSTNAME || process.env.USER || "local",
93+
};
5994

6095
this.acpConnection = createAcpConnection({
6196
adapter: options.adapter,
@@ -74,30 +109,133 @@ export class Agent {
74109
): Promise<void> {
75110
this.logger.info("Attaching PR to task run", { taskId, prUrl, branchName });
76111

77-
if (!this.posthogAPI || !this.taskRunId) {
112+
if (!this.posthogAPI || !this.currentRunId) {
78113
const error = new Error(
79114
"PostHog API not configured or no active run. Cannot attach PR to task.",
80115
);
81116
this.logger.error("PostHog API not configured", error);
82117
throw error;
83118
}
84119

85-
const updates: any = {
120+
const updates: Record<string, unknown> = {
86121
output: { pr_url: prUrl },
87122
};
88123
if (branchName) {
89124
updates.branch = branchName;
90125
}
91126

92-
await this.posthogAPI.updateTaskRun(taskId, this.taskRunId, updates);
127+
await this.posthogAPI.updateTaskRun(taskId, this.currentRunId, updates);
93128
this.logger.debug("PR attached to task run", {
94129
taskId,
95-
taskRunId: this.taskRunId,
130+
taskRunId: this.currentRunId,
96131
prUrl,
97132
});
98133
}
99134

100135
async cleanup(): Promise<void> {
101136
await this.acpConnection?.cleanup();
102137
}
138+
139+
/**
140+
* Stop the agent gracefully, capturing final state.
141+
* Emits a tree_snapshot event with interrupted flag if mid-tool.
142+
*/
143+
async stop(): Promise<TreeSnapshotEvent | null> {
144+
const interrupted = this.isExecutingTool;
145+
146+
this.logger.info("Stopping agent", {
147+
interrupted,
148+
taskId: this.currentTaskId,
149+
runId: this.currentRunId,
150+
});
151+
152+
// Capture tree snapshot immediately
153+
let snapshot: TreeSnapshotEvent | null = null;
154+
if (this.treeTracker) {
155+
const treeSnapshot = await this.treeTracker.captureTree({ interrupted });
156+
if (treeSnapshot) {
157+
snapshot = {
158+
...treeSnapshot,
159+
device: this.deviceInfo,
160+
};
161+
162+
// Emit tree_snapshot event
163+
await this.emitTreeSnapshot(snapshot);
164+
}
165+
}
166+
167+
// Flush session log writer
168+
if (this.sessionLogWriter && this.currentRunId) {
169+
await this.sessionLogWriter.flush(this.currentRunId);
170+
}
171+
172+
this.logger.info("Agent stopped", {
173+
hasSnapshot: !!snapshot,
174+
interrupted,
175+
});
176+
177+
return snapshot;
178+
}
179+
180+
/**
181+
* Emit a tree_snapshot event to the log.
182+
*/
183+
private async emitTreeSnapshot(snapshot: TreeSnapshotEvent): Promise<void> {
184+
if (!this.acpConnection) return;
185+
186+
await this.acpConnection.agentConnection.extNotification?.(
187+
POSTHOG_NOTIFICATIONS.TREE_SNAPSHOT,
188+
snapshot as unknown as Record<string, unknown>,
189+
);
190+
}
191+
192+
/**
193+
* Set device info for tracking where work happens.
194+
*/
195+
setDeviceInfo(info: DeviceInfo): void {
196+
this.deviceInfo = info;
197+
}
198+
199+
/**
200+
* Get current device info.
201+
*/
202+
getDeviceInfo(): DeviceInfo | undefined {
203+
return this.deviceInfo;
204+
}
205+
206+
/**
207+
* Set agent mode (interactive or background).
208+
*/
209+
setAgentMode(mode: AgentMode): void {
210+
const previousMode = this.agentMode;
211+
this.agentMode = mode;
212+
213+
if (previousMode !== mode && this.acpConnection) {
214+
this.acpConnection.agentConnection.extNotification?.(
215+
POSTHOG_NOTIFICATIONS.MODE_CHANGE,
216+
{ mode, previous_mode: previousMode },
217+
);
218+
}
219+
}
220+
221+
/**
222+
* Get current agent mode.
223+
*/
224+
getAgentMode(): AgentMode {
225+
return this.agentMode;
226+
}
227+
228+
/**
229+
* Mark tool execution started/ended (used by stop() to determine interrupted state).
230+
*/
231+
setToolExecuting(executing: boolean): void {
232+
this.isExecutingTool = executing;
233+
}
234+
235+
/**
236+
* Get the tree tracker instance.
237+
*/
238+
getTreeTracker(): TreeTracker | undefined {
239+
return this.treeTracker;
240+
}
103241
}

0 commit comments

Comments
 (0)