Open History to see the last commit
1import { OrchestrationV2TurnItem, RunId, ThreadId, TurnItemId } from "@t3tools/contracts";
2import * as Context from "effect/Context";
3import * as Effect from "effect/Effect";
4import * as Layer from "effect/Layer";
5import * as Schema from "effect/Schema";
6import * as SqlClient from "effect/sql/SqlClient";
7
8export class TurnItemPositionStoreError extends Schema.TaggedError<TurnItemPositionStoreError>()(
9 "TurnItemPositionStoreError",
10 {
11 threadId: ThreadId,
12 turnItemId: TurnItemId,
13 cause: Schema.optional(Schema.Defect()),
14 },
15) {}
16
17const isTurnItemPositionStoreError = Schema.is(TurnItemPositionStoreError);
18
19export interface TurnItemPositionStoreV2Shape {
20 readonly allocate: (input: {
21 readonly threadId: ThreadId;
22 readonly turnItemId: TurnItemId;
23 readonly runId: RunId | null;
24 readonly runOrdinal?: number;
25 // `run_start` takes the run's unused first slot, ahead of its items, and appends once it is taken.
26 readonly placement?: "append" | "run_start";
27 // Keep an existing position even when it lies outside the run's range.
28 readonly keepExisting?: boolean;
29 }) => Effect.Effect<number, TurnItemPositionStoreError>;
30 readonly normalize: (
31 item: OrchestrationV2TurnItem,
32 runOrdinal?: number,
33 ) => Effect.Effect<OrchestrationV2TurnItem, TurnItemPositionStoreError>;
34}
35
36export class TurnItemPositionStoreV2 extends Context.Service<
37 TurnItemPositionStoreV2,
38 TurnItemPositionStoreV2Shape
39>()("t3/orchestration-v2/TurnItemPositionStore/TurnItemPositionStoreV2") {}
40
41export const layer: Layer.Layer<TurnItemPositionStoreV2, never, SqlClient.SqlClient> = Layer.effect(
42 TurnItemPositionStoreV2,
43 Effect.gen(function* () {
44 const sql = yield* SqlClient.SqlClient;
45
46 const allocate: TurnItemPositionStoreV2Shape["allocate"] =
47 /* This helper runs for individual orchestration events; the enclosing operation owns the span. */
48 Effect.fnUntraced(
49 function* ({
50 threadId,
51 turnItemId,
52 runId,
53 runOrdinal: suppliedRunOrdinal,
54 placement = "append",
55 keepExisting = false,
56 }) {
57 const existing = yield* sql<{ readonly ordinal: number }>`
58 SELECT ordinal
59 FROM orchestration_v2_turn_item_positions
60 WHERE thread_id = ${threadId} AND turn_item_id = ${turnItemId}
61 LIMIT 1
62 `;
63 const runRows =
64 runId === null || suppliedRunOrdinal !== undefined
65 ? []
66 : yield* sql<{ readonly ordinal: number }>`
67 SELECT ordinal
68 FROM orchestration_v2_projection_runs
69 WHERE thread_id = ${threadId} AND run_id = ${runId}
70 LIMIT 1
71 `;
72 const runOrdinal = suppliedRunOrdinal ?? runRows[0]?.ordinal ?? null;
73 const lowerBound = runOrdinal === null ? 0 : runOrdinal * 1_000_000;
74 const upperBound = runOrdinal === null ? 999_999 : lowerBound + 999_999;
75 const current = existing[0]?.ordinal;
76 if (current !== undefined) {
77 // A steer that missed its turn starts the next run; it moves into that run's range.
78 if (
79 keepExisting ||
80 runOrdinal === null ||
81 (current >= lowerBound && current <= upperBound)
82 ) {
83 return current;
84 }
85 yield* sql`
86 UPDATE orchestration_v2_turn_item_positions
87 SET ordinal = (
88 SELECT COALESCE(MAX(ordinal), ${lowerBound}) + 1
89 FROM orchestration_v2_turn_item_positions
90 WHERE thread_id = ${threadId}
91 AND ordinal >= ${lowerBound}
92 AND ordinal <= ${upperBound}
93 )
94 WHERE thread_id = ${threadId} AND turn_item_id = ${turnItemId}
95 `;
96 } else {
97 yield* sql`
98 INSERT INTO orchestration_v2_turn_item_positions (thread_id, turn_item_id, ordinal)
99 SELECT
100 ${threadId},
101 ${turnItemId},
102 ${
103 placement === "run_start"
104 ? sql`CASE WHEN COALESCE(MIN(ordinal), ${upperBound}) > ${lowerBound} THEN ${lowerBound} ELSE MAX(ordinal) + 1 END`
105 : sql`COALESCE(MAX(ordinal), ${lowerBound}) + 1`
106 }
107 FROM orchestration_v2_turn_item_positions
108 WHERE thread_id = ${threadId}
109 AND ordinal >= ${lowerBound}
110 AND ordinal <= ${upperBound}
111 ON CONFLICT(thread_id, turn_item_id) DO NOTHING
112 `;
113 }
114 const rows = yield* sql<{ readonly ordinal: number }>`
115 SELECT ordinal
116 FROM orchestration_v2_turn_item_positions
117 WHERE thread_id = ${threadId} AND turn_item_id = ${turnItemId}
118 LIMIT 1
119 `;
120 const ordinal = rows[0]?.ordinal;
121 if (ordinal === undefined) {
122 return yield* new TurnItemPositionStoreError({
123 threadId,
124 turnItemId,
125 cause: "Position allocation returned no row.",
126 });
127 }
128 return ordinal;
129 },
130 (effect, { threadId, turnItemId }) =>
131 effect.pipe(
132 Effect.mapError((cause) =>
133 isTurnItemPositionStoreError(cause)
134 ? cause
135 : new TurnItemPositionStoreError({ threadId, turnItemId, cause }),
136 ),
137 ),
138 );
139
140 return TurnItemPositionStoreV2.of({
141 allocate,
142 // Per-event normalization shares the enclosing orchestration span.
143 normalize: Effect.fnUntraced(function* (item, runOrdinal) {
144 return yield* allocate({
145 threadId: item.threadId,
146 turnItemId: item.id,
147 runId: item.runId,
148 // Queued steers keep their submission order until their run starts (see Orchestrator).
149 keepExisting: item.status === "pending",
150 ...(runOrdinal === undefined ? {} : { runOrdinal }),
151 }).pipe(Effect.map((ordinal) => (item.ordinal === ordinal ? item : { ...item, ordinal })));
152 }),
153 });
154 }),
155);