diff --git a/apps/server/src/persistence/Migrations.ts b/apps/server/src/persistence/Migrations.ts index 024c29a43a636..0d08de141a01e 100644 --- a/apps/server/src/persistence/Migrations.ts +++ b/apps/server/src/persistence/Migrations.ts @@ -91,6 +91,7 @@ import ForkMigration010 from "./Migrations/fork/010_ProjectionThreadParentIndex. import ForkMigration011 from "./Migrations/fork/011_ProfileSnapshots.ts"; import ForkMigration012 from "./Migrations/fork/012_TailnetDevices.ts"; import ForkMigration013 from "./Migrations/fork/013_TailnetDevicePolicy.ts"; +import ForkMigration014 from "./Migrations/fork/014_ThreadTaskDepartures.ts"; const migrationEntries = [ [1, "OrchestrationEvents", Migration0001], @@ -184,6 +185,7 @@ const forkMigrationEntries = [ [11, "ProfileSnapshots", ForkMigration011, 68], [12, "TailnetDevices", ForkMigration012, 72], [13, "TailnetDevicePolicy", ForkMigration013, 72], + [14, "ThreadTaskDepartures", ForkMigration014, 72], ] as const; export const forkMigrationManifest = forkMigrationEntries.map( diff --git a/apps/server/src/persistence/Migrations/fork/014_ThreadTaskDepartures.ts b/apps/server/src/persistence/Migrations/fork/014_ThreadTaskDepartures.ts new file mode 100644 index 0000000000000..442325242a124 --- /dev/null +++ b/apps/server/src/persistence/Migrations/fork/014_ThreadTaskDepartures.ts @@ -0,0 +1,20 @@ +import * as Effect from "effect/Effect"; +import * as SqlClient from "effect/sql/SqlClient"; + +// A task that moved to another owner, under the cursor of that move, so the +// former owner's delta reads drop it (see threadTask/ThreadTaskStore.ts). +export default Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient; + yield* sql` + CREATE TABLE IF NOT EXISTS thread_task_departures ( + worker_thread_id TEXT NOT NULL, + owner_thread_id TEXT NOT NULL, + cursor INTEGER NOT NULL, + PRIMARY KEY (worker_thread_id, owner_thread_id) + ) + `; + yield* sql` + CREATE INDEX IF NOT EXISTS idx_thread_task_departures_owner + ON thread_task_departures(owner_thread_id, cursor) + `; +}); diff --git a/apps/server/src/threadTask/ThreadTaskService.test.ts b/apps/server/src/threadTask/ThreadTaskService.test.ts index a524129e54dab..536cd8a53d216 100644 --- a/apps/server/src/threadTask/ThreadTaskService.test.ts +++ b/apps/server/src/threadTask/ThreadTaskService.test.ts @@ -641,6 +641,80 @@ it.effect("watch returns on the next change and times out without one", () => ).pipe(Effect.provide(testLayer)), ); +it.effect("a read after a cursor returns only changed, settled and departed tasks", () => + Effect.scoped( + Effect.gen(function* () { + const f = yield* fixture; + const store = yield* ThreadTaskStore.ThreadTaskStore; + const tasks = yield* makeService; + yield* tasks.start(); + yield* f.create("captain"); + for (const id of ["alpha", "beta", "gamma"]) { + yield* f.create(id, "captain"); + yield* f.send(id, 1); + } + const alpha = { kind: "thread", threadId: ThreadId.make("alpha") } as const; + yield* tasks.assign(owner, { + threadId: alpha.threadId, + summary: "Alpha", + settleWhenAccepted: true, + }); + yield* tasks.assign(owner, { threadId: ThreadId.make("beta"), summary: "Beta" }); + yield* tasks.assign(owner, { threadId: ThreadId.make("gamma"), summary: "Gamma" }); + const full = yield* tasks.read(owner, {}); + assert.deepEqual(full.tasks.map((view) => view.task.workerThreadId).toSorted(), [ + "alpha", + "beta", + "gamma", + ]); + assert.isUndefined(full.removed); + + yield* tasks.update(owner, { + threadId: ThreadId.make("beta"), + expectedRevision: 1, + summary: "Beta, narrowed", + }); + const changed = yield* tasks.read(owner, { afterCursor: full.cursor }); + assert.deepEqual( + changed.tasks.map((view) => view.task.summary), + ["Beta, narrowed"], + ); + assert.isUndefined(changed.removed); + + yield* tasks.update(alpha, { expectedRevision: 1, status: "DONE", evidence: ["sha 1"] }); + yield* f.endLatestRun("alpha"); + yield* tasks.drain; + const doneCursor = (yield* tasks.read(owner, { afterCursor: changed.cursor })).cursor; + yield* tasks.update(owner, { threadId: alpha.threadId, expectedRevision: 2, accept: true }); + const settled = yield* tasks.read(owner, { afterCursor: doneCursor }); + assert.deepEqual( + settled.tasks.map((view) => [view.task.workerThreadId, view.task.settlement.state]), + [["alpha", "settled"]], + ); + + // A task that moves to another owner leaves the former owner's delta as `removed`. + const gamma = Option.getOrThrow(yield* store.get(ThreadId.make("gamma"))); + yield* store.put({ ...gamma, ownerThreadId: stranger.threadId }); + const departed = yield* tasks.read(owner, { afterCursor: settled.cursor }); + assert.deepEqual(departed.tasks, []); + assert.deepEqual(departed.removed, [ThreadId.make("gamma")]); + assert.deepEqual( + (yield* tasks.read(owner, {})).tasks.map((view) => view.task.workerThreadId).toSorted(), + ["alpha", "beta"], + ); + + // Moved back, it is a change again rather than a departure. + yield* store.put({ ...gamma, ownerThreadId: owner.threadId }); + const returned = yield* tasks.read(owner, { afterCursor: departed.cursor }); + assert.deepEqual( + returned.tasks.map((view) => view.task.workerThreadId), + ["gamma"], + ); + assert.isUndefined(returned.removed); + }), + ).pipe(Effect.provide(testLayer)), +); + it.effect("a captain settles itself after its turn once its children are accepted and quiet", () => Effect.scoped( Effect.gen(function* () { diff --git a/apps/server/src/threadTask/ThreadTaskService.ts b/apps/server/src/threadTask/ThreadTaskService.ts index e21865d19cce8..812226505bd88 100644 --- a/apps/server/src/threadTask/ThreadTaskService.ts +++ b/apps/server/src/threadTask/ThreadTaskService.ts @@ -624,17 +624,38 @@ export const make = Effect.gen(function* () { ), ); - const visibleTasks = (caller: ThreadTaskCaller, afterCursor: number) => - store - .listVisible({ threadId: caller.kind === "user" ? null : caller.threadId, afterCursor }) - .pipe(Effect.catch(unavailable)); + /** Tasks the caller sees that changed after `afterCursor`, and those that left its view. */ + const visibleChanges = (caller: ThreadTaskCaller, afterCursor: number) => + Effect.all({ + tasks: store.listVisible({ + threadId: caller.kind === "user" ? null : caller.threadId, + afterCursor, + }), + departures: + caller.kind === "user" + ? Effect.succeed([]) + : store.listDepartures({ ownerThreadId: caller.threadId, afterCursor }), + }).pipe(Effect.catch(unavailable)); + + const changeResult = Effect.fn("ThreadTaskService.changeResult")(function* ( + changes: Effect.Success>, + cursor: number, + ) { + const removed = changes.departures.map((departure) => departure.workerThreadId); + return { + tasks: yield* Effect.forEach(changes.tasks, view), + ...(removed.length === 0 ? {} : { removed }), + cursor, + timedOut: false, + }; + }); const read: ThreadTaskService["Service"]["read"] = Effect.fn("ThreadTaskService.read")( function* (caller, input) { const cursor = yield* store.latestCursor.pipe(Effect.catch(unavailable)); if (input.threadId === undefined) { - const tasks = yield* visibleTasks(caller, 0); - return { tasks: yield* Effect.forEach(tasks, view), cursor, timedOut: false }; + // The cursor is read first, so a write racing this read is returned again, never lost. + return yield* changeResult(yield* visibleChanges(caller, input.afterCursor ?? 0), cursor); } const task = yield* getTask(input.threadId); if (task === undefined) { @@ -880,22 +901,23 @@ export const make = Effect.gen(function* () { const timeout = Duration.millis( Math.min(input.timeoutMs ?? 30_000, THREAD_TASK_WATCH_MAX_MS), ); - const poll = visibleTasks(caller, afterCursor); + const poll = visibleChanges(caller, afterCursor); + const hasChanges = (changes: Effect.Success) => + changes.tasks.length > 0 || changes.departures.length > 0; const first = yield* poll; - const changed = - first.length > 0 - ? first - : yield* PubSub.take(subscription).pipe( - Effect.andThen(poll), - Effect.repeat({ until: (tasks) => tasks.length > 0 }), - Effect.timeoutOption(timeout), - Effect.map(Option.getOrElse(() => [] as ReadonlyArray)), - ); - return { - tasks: yield* Effect.forEach(changed, view), - cursor: changed.reduce((max, task) => Math.max(max, task.cursor), afterCursor), - timedOut: changed.length === 0, - }; + const changed = hasChanges(first) + ? first + : yield* PubSub.take(subscription).pipe( + Effect.andThen(poll), + Effect.repeat({ until: hasChanges }), + Effect.timeoutOption(timeout), + Effect.map(Option.getOrElse(() => first)), + ); + const cursor = [...changed.tasks, ...changed.departures].reduce( + (max, change) => Math.max(max, change.cursor), + afterCursor, + ); + return { ...(yield* changeResult(changed, cursor)), timedOut: !hasChanges(changed) }; }), ); diff --git a/apps/server/src/threadTask/ThreadTaskStore.ts b/apps/server/src/threadTask/ThreadTaskStore.ts index 56342f85d36b9..1c59932330a60 100644 --- a/apps/server/src/threadTask/ThreadTaskStore.ts +++ b/apps/server/src/threadTask/ThreadTaskStore.ts @@ -73,6 +73,14 @@ export class ThreadTaskStore extends Context.Service< readonly threadId: ThreadId | null; readonly afterCursor: number; }) => Effect.Effect, ThreadTaskStoreError>; + /** Tasks that moved from this owner to another after `afterCursor`, under the move's cursor. */ + readonly listDepartures: (input: { + readonly ownerThreadId: ThreadId; + readonly afterCursor: number; + }) => Effect.Effect< + ReadonlyArray<{ readonly workerThreadId: ThreadId; readonly cursor: number }>, + ThreadTaskStoreError + >; /** Tasks with this project key, newest change first, at most `limit`. */ readonly listByProjectKey: (input: { readonly projectKey: string; @@ -240,6 +248,16 @@ const make = Effect.gen(function* () { ` ).pipe(Effect.mapError(fail("list-visible")), Effect.flatMap(decodeRows("list-visible"))), ), + listDepartures: Effect.fn("ThreadTaskStore.listDepartures")(({ ownerThreadId, afterCursor }) => + sql<{ readonly workerThreadId: ThreadId; readonly cursor: number }>` + SELECT worker_thread_id AS "workerThreadId", cursor FROM thread_task_departures + WHERE owner_thread_id = ${ownerThreadId} AND cursor > ${afterCursor} + ORDER BY cursor + `.pipe( + Effect.map((rows) => rows.map((row) => ({ ...row, cursor: Number(row.cursor) }))), + Effect.mapError(fail("list-departures")), + ), + ), listByProjectKey: Effect.fn("ThreadTaskStore.listByProjectKey")(({ projectKey, limit }) => sql<{ readonly payload: string }>` SELECT payload_json AS payload FROM thread_tasks @@ -269,11 +287,17 @@ const make = Effect.gen(function* () { ), ), latestCursor, - put: Effect.fn("ThreadTaskStore.put")(function* (task) { - const cursor = (yield* latestCursor) + 1; - const stored = { ...task, cursor }; - const payload = yield* encodeTask(stored).pipe(Effect.mapError(fail("encode"))); - yield* sql` + put: Effect.fn("ThreadTaskStore.put")((task) => + Effect.gen(function* () { + const cursor = (yield* latestCursor) + 1; + const stored = { ...task, cursor }; + const payload = yield* encodeTask(stored).pipe(Effect.mapError(fail("encode"))); + const previous = yield* sql<{ readonly owner: ThreadId }>` + SELECT owner_thread_id AS owner FROM thread_tasks + WHERE worker_thread_id = ${stored.workerThreadId} + `.pipe(Effect.mapError(fail("put"))); + const previousOwner = previous[0]?.owner; + yield* sql` INSERT INTO thread_tasks (worker_thread_id, owner_thread_id, cursor, payload_json) VALUES (${stored.workerThreadId}, ${stored.ownerThreadId}, ${cursor}, ${payload}) ON CONFLICT (worker_thread_id) DO UPDATE SET @@ -281,8 +305,25 @@ const make = Effect.gen(function* () { cursor = excluded.cursor, payload_json = excluded.payload_json `.pipe(Effect.mapError(fail("put"))); - return stored; - }), + if (previousOwner !== undefined && previousOwner !== stored.ownerThreadId) { + yield* sql` + INSERT INTO thread_task_departures (worker_thread_id, owner_thread_id, cursor) + VALUES (${stored.workerThreadId}, ${previousOwner}, ${cursor}) + ON CONFLICT (worker_thread_id, owner_thread_id) DO UPDATE SET cursor = excluded.cursor + `.pipe(Effect.mapError(fail("put"))); + // Back with an earlier owner: its delta shows the task again, not a departure. + yield* sql` + DELETE FROM thread_task_departures + WHERE worker_thread_id = ${stored.workerThreadId} + AND owner_thread_id = ${stored.ownerThreadId} + `.pipe(Effect.mapError(fail("put"))); + } + return stored; + }).pipe( + sql.withTransaction, + Effect.catchTags({ SqlError: (cause) => Effect.fail(fail("put")(cause)) }), + ), + ), }); }); diff --git a/packages/contracts/src/threadTask.ts b/packages/contracts/src/threadTask.ts index fecf211a68973..03a337e335bbd 100644 --- a/packages/contracts/src/threadTask.ts +++ b/packages/contracts/src/threadTask.ts @@ -217,6 +217,12 @@ export const ThreadTaskReadInput = Schema.Struct({ "A child chat whose task you own, or your own chat. Omit to list every task you own.", }), ), + afterCursor: Schema.optional( + NonNegativeInt.annotate({ + description: + "Without threadId: only tasks changed after this cursor from an earlier read, plus `removed`. Omit for every task.", + }), + ), }); export type ThreadTaskReadInput = typeof ThreadTaskReadInput.Type; @@ -271,6 +277,8 @@ export type ThreadTaskWatchInput = typeof ThreadTaskWatchInput.Type; export const ThreadTaskListResult = Schema.Struct({ tasks: Schema.Array(ThreadTaskView), + /** Delta reads: worker chats whose tasks moved to another owner; drop them. */ + removed: Schema.optional(Schema.Array(ThreadId)), cursor: NonNegativeInt, timedOut: Schema.Boolean, }); diff --git a/packages/mcp-gateway/src/catalog.ts b/packages/mcp-gateway/src/catalog.ts index 0a9cef216b0cc..a67be279c37be 100644 --- a/packages/mcp-gateway/src/catalog.ts +++ b/packages/mcp-gateway/src/catalog.ts @@ -388,7 +388,7 @@ export const TOOL_SPECS = { { environmentId, ...threadTaskInputFields.assign }, ], t3_task_read: [ - "Read a task: your own as a worker, a child's as its owner, or omit threadId for every task you own or work on. Returns status, revision, needs, evidence, waitingOn, checkBack, whether that continuation is live, the worker's run state, and acceptance and settlement.", + "Read a task: your own as a worker, a child's as its owner, or omit threadId for every task you own or work on. Returns status, revision, needs, evidence, waitingOn, checkBack, whether that continuation is live, the worker's run state, and acceptance and settlement, plus a cursor. To re-check your tasks, pass that cursor as afterCursor (without threadId): you get only tasks changed since, and `removed` lists chats whose tasks moved to another owner.", { environmentId, ...threadTaskInputFields.read }, ], t3_task_update: [ diff --git a/packages/mcp-gateway/src/tools.ts b/packages/mcp-gateway/src/tools.ts index b96290380005b..5a797d109297f 100644 --- a/packages/mcp-gateway/src/tools.ts +++ b/packages/mcp-gateway/src/tools.ts @@ -83,6 +83,14 @@ export const threadTaskInputFields = { .describe( "A child chat whose task you own, or your own chat. Omit to list every task you own.", ), + afterCursor: z + .number() + .int() + .min(0) + .optional() + .describe( + "Without threadId: only tasks changed after this cursor from an earlier read, plus `removed`. Omit for every task.", + ), }, update: { threadId: taskThreadId.optional().describe("The worker chat. Omit to update your own task."),