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 }) => Effect.Effect<number, TurnItemPositionStoreError>;
28 readonly normalize: (
29 item: OrchestrationV2TurnItem,
30 runOrdinal?: number,
31 ) => Effect.Effect<OrchestrationV2TurnItem, TurnItemPositionStoreError>;
32}
33
34export class TurnItemPositionStoreV2 extends Context.Service<
35 TurnItemPositionStoreV2,
36 TurnItemPositionStoreV2Shape
37>()("t3/orchestration-v2/TurnItemPositionStore/TurnItemPositionStoreV2") {}
38
39export const layer: Layer.Layer<TurnItemPositionStoreV2, never, SqlClient.SqlClient> = Layer.effect(
40 TurnItemPositionStoreV2,
41 Effect.gen(function* () {
42 const sql = yield* SqlClient.SqlClient;
43
44 const allocate: TurnItemPositionStoreV2Shape["allocate"] =
45 /* This helper runs for individual orchestration events; the enclosing operation owns the span. */
46 Effect.fnUntraced(
47 function* ({
48 threadId,
49 turnItemId,
50 runId,
51 runOrdinal: suppliedRunOrdinal,
52 placement = "append",
53 }) {
54 const existing = yield* sql<{ readonly ordinal: number }>`
55 SELECT ordinal
56 FROM orchestration_v2_turn_item_positions
57 WHERE thread_id = ${threadId} AND turn_item_id = ${turnItemId}
58 LIMIT 1
59 `;
60 if (existing[0] !== undefined) return existing[0].ordinal;
61
62 const runRows =
63 runId === null
64 ? []
65 : yield* sql<{ readonly ordinal: number }>`
66 SELECT ordinal
67 FROM orchestration_v2_projection_runs
68 WHERE thread_id = ${threadId} AND run_id = ${runId}
69 LIMIT 1
70 `;
71 const runOrdinal = suppliedRunOrdinal ?? runRows[0]?.ordinal ?? null;
72 const lowerBound = runOrdinal === null ? 0 : runOrdinal * 1_000_000;
73 const upperBound = runOrdinal === null ? 999_999 : lowerBound + 999_999;
74 yield* sql`
75 INSERT INTO orchestration_v2_turn_item_positions (thread_id, turn_item_id, ordinal)
76 SELECT
77 ${threadId},
78 ${turnItemId},
79 ${
80 placement === "run_start"
81 ? sql`CASE WHEN COALESCE(MIN(ordinal), ${upperBound}) > ${lowerBound} THEN ${lowerBound} ELSE MAX(ordinal) + 1 END`
82 : sql`COALESCE(MAX(ordinal), ${lowerBound}) + 1`
83 }
84 FROM orchestration_v2_turn_item_positions
85 WHERE thread_id = ${threadId}
86 AND ordinal >= ${lowerBound}
87 AND ordinal <= ${upperBound}
88 ON CONFLICT(thread_id, turn_item_id) DO NOTHING
89 `;
90 const rows = yield* sql<{ readonly ordinal: number }>`
91 SELECT ordinal
92 FROM orchestration_v2_turn_item_positions
93 WHERE thread_id = ${threadId} AND turn_item_id = ${turnItemId}
94 LIMIT 1
95 `;
96 const ordinal = rows[0]?.ordinal;
97 if (ordinal === undefined) {
98 return yield* new TurnItemPositionStoreError({
99 threadId,
100 turnItemId,
101 cause: "Position allocation returned no row.",
102 });
103 }
104 return ordinal;
105 },
106 (effect, { threadId, turnItemId }) =>
107 effect.pipe(
108 Effect.mapError((cause) =>
109 isTurnItemPositionStoreError(cause)
110 ? cause
111 : new TurnItemPositionStoreError({ threadId, turnItemId, cause }),
112 ),
113 ),
114 );
115
116 return TurnItemPositionStoreV2.of({
117 allocate,
118 // Per-event normalization shares the enclosing orchestration span.
119 normalize: Effect.fnUntraced(function* (item, runOrdinal) {
120 return yield* allocate({
121 threadId: item.threadId,
122 turnItemId: item.id,
123 runId: item.runId,
124 ...(runOrdinal === undefined ? {} : { runOrdinal }),
125 }).pipe(Effect.map((ordinal) => (item.ordinal === ordinal ? item : { ...item, ordinal })));
126 }),
127 });
128 }),
129);