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