Skip to content

Commit 4986279

Browse files
RealYusufIsmailRealYusufIsmail
committed
fix(server): validate custom agents before direct thread creation
1 parent 24a448e commit 4986279

3 files changed

Lines changed: 249 additions & 26 deletions

File tree

‎apps/server/src/orchestration-v2/ThreadLaunchService.test.ts‎

Lines changed: 44 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -128,11 +128,6 @@ function makeHarness(options: HarnessOptions = {}) {
128128
);
129129
const layerServerSettings = ServerSettings.layerTest(options.serverSettings);
130130
const layerReceipts = CommandReceiptStore.layer.pipe(Layer.provide(layerDatabase));
131-
const layerThreadManagement = ThreadManagement.layer.pipe(
132-
Layer.provide(layerOrchestrator),
133-
Layer.provide(layerServerSettings),
134-
Layer.provide(layerReceipts),
135-
);
136131
const layerOutbox = EffectOutbox.layer.pipe(Layer.provide(layerDatabase));
137132
const createWorktree = vi.fn(
138133
options.createWorktree ??
@@ -157,30 +152,49 @@ function makeHarness(options: HarnessOptions = {}) {
157152
const generateThreadTitle = vi.fn(
158153
options.generateTitle ?? (() => Effect.succeed({ title: "Generated title" })),
159154
);
155+
const layerProjectService = Layer.succeed(ProjectService.ProjectService, {
156+
create: () => Effect.die("unused"),
157+
bootstrap: () => Effect.die("unused"),
158+
update: () => Effect.die("unused"),
159+
delete: () => Effect.die("unused"),
160+
getById: (id) =>
161+
Effect.succeed(
162+
id === projectId
163+
? Option.some(configuredProject)
164+
: id === otherProjectId
165+
? Option.some(otherProject)
166+
: Option.none(),
167+
),
168+
getByWorkspaceRoot: () => Effect.succeed(Option.some(configuredProject)),
169+
snapshot: Effect.die("unused"),
170+
getShell: () => Effect.die("unused"),
171+
listShells: () => Effect.die("unused"),
172+
listAgents: () =>
173+
Effect.succeed([
174+
{
175+
id: AgentId.make("lead"),
176+
name: "Lead",
177+
description: "Coordinates work.",
178+
},
179+
{
180+
id: AgentId.make("reasoner"),
181+
name: "Reasoner",
182+
description: "Reasons carefully.",
183+
tier: "reasoning",
184+
},
185+
{
186+
id: AgentId.make("inheritor"),
187+
name: "Inheritor",
188+
description: "Inherits the selected model.",
189+
},
190+
]),
191+
});
160192
const layerExternalServices = Layer.mergeAll(
161193
NodeServices.layer,
162194
WorktreeSetupTracker.layer,
163195
Layer.mock(ProjectCloneTracker.ProjectCloneTracker)({ get: () => Effect.succeed(null) }),
164196
Layer.mock(TerminalManager.TerminalManager)({ close: () => Effect.void }),
165-
Layer.succeed(ProjectService.ProjectService, {
166-
create: () => Effect.die("unused"),
167-
bootstrap: () => Effect.die("unused"),
168-
update: () => Effect.die("unused"),
169-
delete: () => Effect.die("unused"),
170-
getById: (id) =>
171-
Effect.succeed(
172-
id === projectId
173-
? Option.some(configuredProject)
174-
: id === otherProjectId
175-
? Option.some(otherProject)
176-
: Option.none(),
177-
),
178-
getByWorkspaceRoot: () => Effect.succeed(Option.some(configuredProject)),
179-
snapshot: Effect.die("unused"),
180-
getShell: () => Effect.die("unused"),
181-
listShells: () => Effect.die("unused"),
182-
listAgents: () => Effect.die("unused"),
183-
}),
197+
layerProjectService,
184198
Layer.mock(GitWorkflow.GitWorkflowService)({
185199
createWorktree,
186200
renameBranch,
@@ -208,6 +222,12 @@ function makeHarness(options: HarnessOptions = {}) {
208222
folderForThread: () => Effect.succeed(Option.none()),
209223
}),
210224
);
225+
const layerThreadManagement = ThreadManagement.layer.pipe(
226+
Layer.provide(layerOrchestrator),
227+
Layer.provide(layerServerSettings),
228+
Layer.provide(layerReceipts),
229+
Layer.provide(layerProjectService),
230+
);
211231
const layerLaunch = ThreadLaunch.layer.pipe(
212232
Layer.provide(
213233
Layer.mergeAll(
@@ -2051,7 +2071,6 @@ it.effect("uses the persisted thread model when resuming after an accepted creat
20512071
projectId: input.projectId,
20522072
title: input.title,
20532073
modelSelection: persistedModelSelection,
2054-
agentId: AgentId.make("lead"),
20552074
runtimeMode: input.runtimeMode,
20562075
interactionMode: input.interactionMode,
20572076
branch: null,

‎apps/server/src/orchestration-v2/ThreadManagementService.test.ts‎

Lines changed: 142 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { expect, it } from "@effect/vitest";
22
import {
3+
AgentId,
34
CommandId,
45
MessageId,
56
NodeId,
@@ -32,6 +33,7 @@ import * as TestClock from "effect/testing/TestClock";
3233

3334
import * as LegacyV1ThreadImporter from "./legacy/LegacyV1ThreadImporter.ts";
3435
import * as ServerSettings from "../serverSettings.ts";
36+
import * as ProjectService from "../project/ProjectService.ts";
3537
import * as SqlitePersistence from "../persistence/Sqlite.ts";
3638
import * as CommandReceiptStore from "./CommandReceiptStore.ts";
3739
import * as Orchestrator from "./Orchestrator.ts";
@@ -359,7 +361,8 @@ it.effect("inherits the machine value when the project has no override", () => {
359361
});
360362

361363
const layerDispatchChokepointTest =
362-
(settingsLayer: ReturnType<typeof ServerSettings.layerTest>) => (dispatched: Array<string>) =>
364+
(settingsLayer: ReturnType<typeof ServerSettings.layerTest>) =>
365+
(dispatched: Array<string>, agentListCalls: Array<ProjectId> = [], validAgentId?: AgentId) =>
363366
ThreadManagementService.layer.pipe(
364367
Layer.provide(Layer.mock(ProviderSessionManager.ProviderSessionManagerV2)({})),
365368
Layer.provide(
@@ -383,6 +386,17 @@ const layerDispatchChokepointTest =
383386
}),
384387
}),
385388
),
389+
Layer.provideMerge(
390+
Layer.mock(ProjectService.ProjectService)({
391+
listAgents: ({ projectId }) =>
392+
Effect.sync(() => {
393+
agentListCalls.push(projectId);
394+
return validAgentId === undefined
395+
? []
396+
: [{ id: validAgentId, name: "Selected agent", description: "Test agent." }];
397+
}),
398+
}),
399+
),
386400
Layer.provide(settingsLayer),
387401
Layer.provideMerge(layerReceiptsTest),
388402
);
@@ -476,6 +490,133 @@ it.effect(
476490
},
477491
);
478492

493+
it.effect("rejects a direct thread.create with an unknown custom agent before dispatch", () => {
494+
const projectId = ProjectId.make("project:thread-management:unknown-agent");
495+
const threadId = ThreadId.make("thread:thread-management:unknown-agent");
496+
const commandId = CommandId.make("command:thread-management:unknown-agent");
497+
const dispatched: Array<string> = [];
498+
const agentListCalls: Array<ProjectId> = [];
499+
const layerTest = layerDispatchChokepointTest(ServerSettings.layerTest())(
500+
dispatched,
501+
agentListCalls,
502+
);
503+
504+
return Effect.gen(function* () {
505+
const threads = yield* ThreadManagementService.ThreadManagementService;
506+
const error = yield* Effect.flip(
507+
threads.dispatch({
508+
type: "thread.create",
509+
commandId,
510+
threadId,
511+
projectId,
512+
title: "Unknown agent",
513+
modelSelection: {
514+
instanceId: ProviderInstanceId.make("codex"),
515+
model: "test",
516+
},
517+
agentId: AgentId.make("missing-agent"),
518+
runtimeMode: "full-access",
519+
interactionMode: "default",
520+
branch: null,
521+
worktreePath: null,
522+
createdBy: "user",
523+
creationSource: "web",
524+
}),
525+
);
526+
527+
expect(error).toBeInstanceOf(Orchestrator.OrchestratorCommandRejectedError);
528+
expect(dispatched).toEqual([]);
529+
expect(agentListCalls).toEqual([projectId]);
530+
}).pipe(Effect.provide(layerTest));
531+
});
532+
533+
it.effect("accepts a direct thread.create when its custom agent exists", () => {
534+
const projectId = ProjectId.make("project:thread-management:known-agent");
535+
const threadId = ThreadId.make("thread:thread-management:known-agent");
536+
const commandId = CommandId.make("command:thread-management:known-agent");
537+
const agentId = AgentId.make("selected-agent");
538+
const dispatched: Array<string> = [];
539+
const agentListCalls: Array<ProjectId> = [];
540+
const layerTest = layerDispatchChokepointTest(ServerSettings.layerTest())(
541+
dispatched,
542+
agentListCalls,
543+
agentId,
544+
);
545+
546+
return Effect.gen(function* () {
547+
const threads = yield* ThreadManagementService.ThreadManagementService;
548+
yield* threads.dispatch({
549+
type: "thread.create",
550+
commandId,
551+
threadId,
552+
projectId,
553+
title: "Known agent",
554+
modelSelection: {
555+
instanceId: ProviderInstanceId.make("codex"),
556+
model: "test",
557+
},
558+
agentId,
559+
runtimeMode: "full-access",
560+
interactionMode: "default",
561+
branch: null,
562+
worktreePath: null,
563+
createdBy: "user",
564+
creationSource: "web",
565+
});
566+
567+
expect(dispatched).toEqual(["thread.create"]);
568+
expect(agentListCalls).toEqual([projectId]);
569+
}).pipe(Effect.provide(layerTest));
570+
});
571+
572+
it.effect("replays an accepted thread.create without revalidating its custom agent", () => {
573+
const projectId = ProjectId.make("project:thread-management:agent-replay");
574+
const threadId = ThreadId.make("thread:thread-management:agent-replay");
575+
const commandId = CommandId.make("command:thread-management:agent-replay");
576+
const dispatched: Array<string> = [];
577+
const agentListCalls: Array<ProjectId> = [];
578+
const layerTest = layerDispatchChokepointTest(ServerSettings.layerTest())(
579+
dispatched,
580+
agentListCalls,
581+
);
582+
583+
return Effect.gen(function* () {
584+
const receipts = yield* CommandReceiptStore.CommandReceiptStoreV2;
585+
yield* receipts.insertIfAbsent({
586+
commandId,
587+
threadId,
588+
commandType: "thread.create",
589+
acceptedAt: yield* DateTime.now,
590+
resultSequence: 1,
591+
status: "accepted",
592+
error: null,
593+
});
594+
595+
const threads = yield* ThreadManagementService.ThreadManagementService;
596+
yield* threads.dispatch({
597+
type: "thread.create",
598+
commandId,
599+
threadId,
600+
projectId,
601+
title: "Accepted agent replay",
602+
modelSelection: {
603+
instanceId: ProviderInstanceId.make("codex"),
604+
model: "test",
605+
},
606+
agentId: AgentId.make("removed-agent"),
607+
runtimeMode: "full-access",
608+
interactionMode: "default",
609+
branch: null,
610+
worktreePath: null,
611+
createdBy: "user",
612+
creationSource: "web",
613+
});
614+
615+
expect(dispatched).toEqual(["thread.create"]);
616+
expect(agentListCalls).toEqual([]);
617+
}).pipe(Effect.provide(layerTest));
618+
});
619+
479620
it.effect("classifies projection infrastructure failures separately from a missing thread", () => {
480621
const projectId = ProjectId.make("project:thread-management:projection-failure");
481622
const threadId = ThreadId.make("thread:thread-management:projection-failure");

‎apps/server/src/orchestration-v2/ThreadManagementService.ts‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts";
5050
import { DispatchModeLimit, exceededDispatchModeLimit } from "./DispatchModeLimit.ts";
5151
import { projectTurnItemForDetail } from "./WireProjection.ts";
5252
import * as LegacyV1ThreadImporter from "./legacy/LegacyV1ThreadImporter.ts";
53+
import * as ProjectService from "../project/ProjectService.ts";
5354
import * as ServerSettings from "../serverSettings.ts";
5455

5556
export type ThreadManagementSendMode = "auto" | "queue" | "steer" | "restart";
@@ -652,9 +653,71 @@ const make = Effect.gen(function* () {
652653
}
653654
};
654655

656+
const validateThreadCreateAgent = (command: OrchestrationV2ServerCommand) => {
657+
if (command.type !== "thread.create" || command.agentId === undefined) {
658+
return Effect.void;
659+
}
660+
661+
return receipts.getByCommandId(command.commandId).pipe(
662+
Effect.mapError(
663+
(cause) =>
664+
new Orchestrator.OrchestratorDispatchError({
665+
commandId: command.commandId,
666+
commandType: command.type,
667+
cause,
668+
}),
669+
),
670+
Effect.flatMap((receipt) => {
671+
// Preserve receipt replay and conflict semantics owned by the orchestrator.
672+
if (Option.isSome(receipt)) {
673+
return Effect.void;
674+
}
675+
676+
return Effect.serviceOption(ProjectService.ProjectService).pipe(
677+
Effect.flatMap((projects) => {
678+
if (Option.isNone(projects)) {
679+
return Effect.fail(
680+
new Orchestrator.OrchestratorDispatchError({
681+
commandId: command.commandId,
682+
commandType: command.type,
683+
cause: new Error("Project service is unavailable for custom agent validation."),
684+
}),
685+
);
686+
}
687+
688+
return projects.value.listAgents({ projectId: command.projectId }).pipe(
689+
Effect.mapError(
690+
(cause) =>
691+
new Orchestrator.OrchestratorDispatchError({
692+
commandId: command.commandId,
693+
commandType: command.type,
694+
cause,
695+
}),
696+
),
697+
Effect.flatMap((agents) =>
698+
agents.some((agent) => agent.id === command.agentId)
699+
? Effect.void
700+
: Effect.fail(
701+
new Orchestrator.OrchestratorCommandRejectedError({
702+
commandId: command.commandId,
703+
commandType: command.type,
704+
cause: new Error(
705+
`Custom agent '${command.agentId}' was not found in the project workspace.`,
706+
),
707+
}),
708+
),
709+
),
710+
);
711+
}),
712+
);
713+
}),
714+
);
715+
};
716+
655717
const dispatch: ThreadManagementServiceShape["dispatch"] = (command) =>
656718
ensureCommandTranscripts(command).pipe(
657719
Effect.andThen(assertDispatchTargetEnabled(command)),
720+
Effect.andThen(validateThreadCreateAgent(command)),
658721
Effect.andThen(orchestrator.dispatch(command)),
659722
);
660723

0 commit comments

Comments
 (0)