Skip to content

Commit 55a2d12

Browse files
steipetesteipete
committed
refactor: split inbound and reload pipelines into staged modules
1 parent 99a3db6 commit 55a2d12

9 files changed

Lines changed: 790 additions & 471 deletions

File tree

‎src/agents/session-tool-result-guard.ts‎

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
HARD_MAX_TOOL_RESULT_CHARS,
1010
truncateToolResultMessage,
1111
} from "./pi-embedded-runner/tool-result-truncation.js";
12+
import { createPendingToolCallState } from "./session-tool-result-state.js";
1213
import { makeMissingToolResult, sanitizeToolCallInputs } from "./session-transcript-repair.js";
1314
import { extractToolCallsFromAssistant, extractToolResultId } from "./tool-call-id.js";
1415

@@ -106,7 +107,7 @@ export function installSessionToolResultGuard(
106107
getPendingIds: () => string[];
107108
} {
108109
const originalAppend = sessionManager.appendMessage.bind(sessionManager);
109-
const pending = new Map<string, string | undefined>();
110+
const pendingState = createPendingToolCallState();
110111
const persistMessage = (message: AgentMessage) => {
111112
const transformer = opts?.transformMessageForPersistence;
112113
return transformer ? transformer(message) : message;
@@ -142,11 +143,11 @@ export function installSessionToolResultGuard(
142143
};
143144

144145
const flushPendingToolResults = () => {
145-
if (pending.size === 0) {
146+
if (pendingState.size() === 0) {
146147
return;
147148
}
148149
if (allowSyntheticToolResults) {
149-
for (const [id, name] of pending.entries()) {
150+
for (const [id, name] of pendingState.entries()) {
150151
const synthetic = makeMissingToolResult({ toolCallId: id, toolName: name });
151152
const flushed = applyBeforeWriteHook(
152153
persistToolResult(persistMessage(synthetic), {
@@ -160,7 +161,7 @@ export function installSessionToolResultGuard(
160161
}
161162
}
162163
}
163-
pending.clear();
164+
pendingState.clear();
164165
};
165166

166167
const guardedAppend = (message: AgentMessage) => {
@@ -171,7 +172,7 @@ export function installSessionToolResultGuard(
171172
allowedToolNames: opts?.allowedToolNames,
172173
});
173174
if (sanitized.length === 0) {
174-
if (pending.size > 0) {
175+
if (pendingState.shouldFlushForSanitizedDrop()) {
175176
flushPendingToolResults();
176177
}
177178
return undefined;
@@ -182,9 +183,9 @@ export function installSessionToolResultGuard(
182183

183184
if (nextRole === "toolResult") {
184185
const id = extractToolResultId(nextMessage as Extract<AgentMessage, { role: "toolResult" }>);
185-
const toolName = id ? pending.get(id) : undefined;
186+
const toolName = id ? pendingState.getToolName(id) : undefined;
186187
if (id) {
187-
pending.delete(id);
188+
pendingState.delete(id);
188189
}
189190
const normalizedToolResult = normalizePersistedToolResultName(nextMessage, toolName);
190191
// Apply hard size cap before persistence to prevent oversized tool results
@@ -221,11 +222,11 @@ export function installSessionToolResultGuard(
221222
// synthetic results (e.g. OpenAI) accumulate stale pending state when a user message
222223
// interrupts in-flight tool calls, leaving orphaned tool_use blocks in the transcript
223224
// that cause API 400 errors on subsequent requests.
224-
if (pending.size > 0 && (toolCalls.length === 0 || nextRole !== "assistant")) {
225+
if (pendingState.shouldFlushBeforeNonToolResult(nextRole, toolCalls.length)) {
225226
flushPendingToolResults();
226227
}
227228
// If new tool calls arrive while older ones are pending, flush the old ones first.
228-
if (pending.size > 0 && toolCalls.length > 0) {
229+
if (pendingState.shouldFlushBeforeNewToolCalls(toolCalls.length)) {
229230
flushPendingToolResults();
230231
}
231232

@@ -243,9 +244,7 @@ export function installSessionToolResultGuard(
243244
}
244245

245246
if (toolCalls.length > 0) {
246-
for (const call of toolCalls) {
247-
pending.set(call.id, call.name);
248-
}
247+
pendingState.trackToolCalls(toolCalls);
249248
}
250249

251250
return result;
@@ -256,6 +255,6 @@ export function installSessionToolResultGuard(
256255

257256
return {
258257
flushPendingToolResults,
259-
getPendingIds: () => Array.from(pending.keys()),
258+
getPendingIds: pendingState.getPendingIds,
260259
};
261260
}

‎src/agents/session-tool-result-state.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
export type PendingToolCall = { id: string; name?: string };
2+
3+
export type PendingToolCallState = {
4+
size: () => number;
5+
entries: () => IterableIterator<[string, string | undefined]>;
6+
getToolName: (id: string) => string | undefined;
7+
delete: (id: string) => void;
8+
clear: () => void;
9+
trackToolCalls: (calls: PendingToolCall[]) => void;
10+
getPendingIds: () => string[];
11+
shouldFlushForSanitizedDrop: () => boolean;
12+
shouldFlushBeforeNonToolResult: (nextRole: unknown, toolCallCount: number) => boolean;
13+
shouldFlushBeforeNewToolCalls: (toolCallCount: number) => boolean;
14+
};
15+
16+
export function createPendingToolCallState(): PendingToolCallState {
17+
const pending = new Map<string, string | undefined>();
18+
19+
return {
20+
size: () => pending.size,
21+
entries: () => pending.entries(),
22+
getToolName: (id: string) => pending.get(id),
23+
delete: (id: string) => {
24+
pending.delete(id);
25+
},
26+
clear: () => {
27+
pending.clear();
28+
},
29+
trackToolCalls: (calls: PendingToolCall[]) => {
30+
for (const call of calls) {
31+
pending.set(call.id, call.name);
32+
}
33+
},
34+
getPendingIds: () => Array.from(pending.keys()),
35+
shouldFlushForSanitizedDrop: () => pending.size > 0,
36+
shouldFlushBeforeNonToolResult: (nextRole: unknown, toolCallCount: number) =>
37+
pending.size > 0 && (toolCallCount === 0 || nextRole !== "assistant"),
38+
shouldFlushBeforeNewToolCalls: (toolCallCount: number) => pending.size > 0 && toolCallCount > 0,
39+
};
40+
}

‎src/gateway/config-reload-plan.ts‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
import { type ChannelId, listChannelPlugins } from "../channels/plugins/index.js";
2+
import { getActivePluginRegistry } from "../plugins/runtime.js";
3+
4+
export type ChannelKind = ChannelId;
5+
6+
export type GatewayReloadPlan = {
7+
changedPaths: string[];
8+
restartGateway: boolean;
9+
restartReasons: string[];
10+
hotReasons: string[];
11+
reloadHooks: boolean;
12+
restartGmailWatcher: boolean;
13+
restartBrowserControl: boolean;
14+
restartCron: boolean;
15+
restartHeartbeat: boolean;
16+
restartHealthMonitor: boolean;
17+
restartChannels: Set<ChannelKind>;
18+
noopPaths: string[];
19+
};
20+
21+
type ReloadRule = {
22+
prefix: string;
23+
kind: "restart" | "hot" | "none";
24+
actions?: ReloadAction[];
25+
};
26+
27+
type ReloadAction =
28+
| "reload-hooks"
29+
| "restart-gmail-watcher"
30+
| "restart-browser-control"
31+
| "restart-cron"
32+
| "restart-heartbeat"
33+
| "restart-health-monitor"
34+
| `restart-channel:${ChannelId}`;
35+
36+
const BASE_RELOAD_RULES: ReloadRule[] = [
37+
{ prefix: "gateway.remote", kind: "none" },
38+
{ prefix: "gateway.reload", kind: "none" },
39+
{
40+
prefix: "gateway.channelHealthCheckMinutes",
41+
kind: "hot",
42+
actions: ["restart-health-monitor"],
43+
},
44+
// Stuck-session warning threshold is read by the diagnostics heartbeat loop.
45+
{ prefix: "diagnostics.stuckSessionWarnMs", kind: "none" },
46+
{ prefix: "hooks.gmail", kind: "hot", actions: ["restart-gmail-watcher"] },
47+
{ prefix: "hooks", kind: "hot", actions: ["reload-hooks"] },
48+
{
49+
prefix: "agents.defaults.heartbeat",
50+
kind: "hot",
51+
actions: ["restart-heartbeat"],
52+
},
53+
{
54+
prefix: "agents.defaults.model",
55+
kind: "hot",
56+
actions: ["restart-heartbeat"],
57+
},
58+
{
59+
prefix: "models",
60+
kind: "hot",
61+
actions: ["restart-heartbeat"],
62+
},
63+
{ prefix: "agent.heartbeat", kind: "hot", actions: ["restart-heartbeat"] },
64+
{ prefix: "cron", kind: "hot", actions: ["restart-cron"] },
65+
{
66+
prefix: "browser",
67+
kind: "hot",
68+
actions: ["restart-browser-control"],
69+
},
70+
];
71+
72+
const BASE_RELOAD_RULES_TAIL: ReloadRule[] = [
73+
{ prefix: "meta", kind: "none" },
74+
{ prefix: "identity", kind: "none" },
75+
{ prefix: "wizard", kind: "none" },
76+
{ prefix: "logging", kind: "none" },
77+
{ prefix: "agents", kind: "none" },
78+
{ prefix: "tools", kind: "none" },
79+
{ prefix: "bindings", kind: "none" },
80+
{ prefix: "audio", kind: "none" },
81+
{ prefix: "agent", kind: "none" },
82+
{ prefix: "routing", kind: "none" },
83+
{ prefix: "messages", kind: "none" },
84+
{ prefix: "session", kind: "none" },
85+
{ prefix: "talk", kind: "none" },
86+
{ prefix: "skills", kind: "none" },
87+
{ prefix: "secrets", kind: "none" },
88+
{ prefix: "plugins", kind: "restart" },
89+
{ prefix: "ui", kind: "none" },
90+
{ prefix: "gateway", kind: "restart" },
91+
{ prefix: "discovery", kind: "restart" },
92+
{ prefix: "canvasHost", kind: "restart" },
93+
];
94+
95+
let cachedReloadRules: ReloadRule[] | null = null;
96+
let cachedRegistry: ReturnType<typeof getActivePluginRegistry> | null = null;
97+
98+
function listReloadRules(): ReloadRule[] {
99+
const registry = getActivePluginRegistry();
100+
if (registry !== cachedRegistry) {
101+
cachedReloadRules = null;
102+
cachedRegistry = registry;
103+
}
104+
if (cachedReloadRules) {
105+
return cachedReloadRules;
106+
}
107+
// Channel docking: plugins contribute hot reload/no-op prefixes here.
108+
const channelReloadRules: ReloadRule[] = listChannelPlugins().flatMap((plugin) => [
109+
...(plugin.reload?.configPrefixes ?? []).map(
110+
(prefix): ReloadRule => ({
111+
prefix,
112+
kind: "hot",
113+
actions: [`restart-channel:${plugin.id}` as ReloadAction],
114+
}),
115+
),
116+
...(plugin.reload?.noopPrefixes ?? []).map(
117+
(prefix): ReloadRule => ({
118+
prefix,
119+
kind: "none",
120+
}),
121+
),
122+
]);
123+
const rules = [...BASE_RELOAD_RULES, ...channelReloadRules, ...BASE_RELOAD_RULES_TAIL];
124+
cachedReloadRules = rules;
125+
return rules;
126+
}
127+
128+
function matchRule(path: string): ReloadRule | null {
129+
for (const rule of listReloadRules()) {
130+
if (path === rule.prefix || path.startsWith(`${rule.prefix}.`)) {
131+
return rule;
132+
}
133+
}
134+
return null;
135+
}
136+
137+
export function buildGatewayReloadPlan(changedPaths: string[]): GatewayReloadPlan {
138+
const plan: GatewayReloadPlan = {
139+
changedPaths,
140+
restartGateway: false,
141+
restartReasons: [],
142+
hotReasons: [],
143+
reloadHooks: false,
144+
restartGmailWatcher: false,
145+
restartBrowserControl: false,
146+
restartCron: false,
147+
restartHeartbeat: false,
148+
restartHealthMonitor: false,
149+
restartChannels: new Set(),
150+
noopPaths: [],
151+
};
152+
153+
const applyAction = (action: ReloadAction) => {
154+
if (action.startsWith("restart-channel:")) {
155+
const channel = action.slice("restart-channel:".length) as ChannelId;
156+
plan.restartChannels.add(channel);
157+
return;
158+
}
159+
switch (action) {
160+
case "reload-hooks":
161+
plan.reloadHooks = true;
162+
break;
163+
case "restart-gmail-watcher":
164+
plan.restartGmailWatcher = true;
165+
break;
166+
case "restart-browser-control":
167+
plan.restartBrowserControl = true;
168+
break;
169+
case "restart-cron":
170+
plan.restartCron = true;
171+
break;
172+
case "restart-heartbeat":
173+
plan.restartHeartbeat = true;
174+
break;
175+
case "restart-health-monitor":
176+
plan.restartHealthMonitor = true;
177+
break;
178+
default:
179+
break;
180+
}
181+
};
182+
183+
for (const path of changedPaths) {
184+
const rule = matchRule(path);
185+
if (!rule) {
186+
plan.restartGateway = true;
187+
plan.restartReasons.push(path);
188+
continue;
189+
}
190+
if (rule.kind === "restart") {
191+
plan.restartGateway = true;
192+
plan.restartReasons.push(path);
193+
continue;
194+
}
195+
if (rule.kind === "none") {
196+
plan.noopPaths.push(path);
197+
continue;
198+
}
199+
plan.hotReasons.push(path);
200+
for (const action of rule.actions ?? []) {
201+
applyAction(action);
202+
}
203+
}
204+
205+
if (plan.restartGmailWatcher) {
206+
plan.reloadHooks = true;
207+
}
208+
209+
return plan;
210+
}

‎src/gateway/config-reload.test.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,53 @@ describe("buildGatewayReloadPlan", () => {
188188
const plan = buildGatewayReloadPlan(["unknownField"]);
189189
expect(plan.restartGateway).toBe(true);
190190
});
191+
192+
it.each([
193+
{
194+
path: "gateway.channelHealthCheckMinutes",
195+
expectRestartGateway: false,
196+
expectHotPath: "gateway.channelHealthCheckMinutes",
197+
expectRestartHealthMonitor: true,
198+
},
199+
{
200+
path: "hooks.gmail.account",
201+
expectRestartGateway: false,
202+
expectHotPath: "hooks.gmail.account",
203+
expectRestartGmailWatcher: true,
204+
expectReloadHooks: true,
205+
},
206+
{
207+
path: "gateway.remote.url",
208+
expectRestartGateway: false,
209+
expectNoopPath: "gateway.remote.url",
210+
},
211+
{
212+
path: "unknownField",
213+
expectRestartGateway: true,
214+
expectRestartReason: "unknownField",
215+
},
216+
])("classifies reload path: $path", (testCase) => {
217+
const plan = buildGatewayReloadPlan([testCase.path]);
218+
expect(plan.restartGateway).toBe(testCase.expectRestartGateway);
219+
if (testCase.expectHotPath) {
220+
expect(plan.hotReasons).toContain(testCase.expectHotPath);
221+
}
222+
if (testCase.expectNoopPath) {
223+
expect(plan.noopPaths).toContain(testCase.expectNoopPath);
224+
}
225+
if (testCase.expectRestartReason) {
226+
expect(plan.restartReasons).toContain(testCase.expectRestartReason);
227+
}
228+
if (testCase.expectRestartHealthMonitor) {
229+
expect(plan.restartHealthMonitor).toBe(true);
230+
}
231+
if (testCase.expectRestartGmailWatcher) {
232+
expect(plan.restartGmailWatcher).toBe(true);
233+
}
234+
if (testCase.expectReloadHooks) {
235+
expect(plan.reloadHooks).toBe(true);
236+
}
237+
});
191238
});
192239

193240
describe("resolveGatewayReloadSettings", () => {

0 commit comments

Comments
 (0)