@sixb/pg 0.1.6 → 0.1.7
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/agents/rows.d.ts +8 -5
- package/dist/agents/rows.d.ts.map +1 -1
- package/dist/agents/runs.d.ts +3 -2
- package/dist/agents/runs.d.ts.map +1 -1
- package/dist/agents/threads.d.ts.map +1 -1
- package/dist/ai-cost-groups.d.ts +5 -0
- package/dist/ai-cost-groups.d.ts.map +1 -0
- package/dist/auth-storage/auth-storage.d.ts +4 -1
- package/dist/auth-storage/auth-storage.d.ts.map +1 -1
- package/dist/auth-storage/device-authorizations.d.ts +29 -0
- package/dist/auth-storage/device-authorizations.d.ts.map +1 -0
- package/dist/auth-storage/rows.d.ts +19 -1
- package/dist/auth-storage/rows.d.ts.map +1 -1
- package/dist/auth-storage/shared.d.ts +1 -1
- package/dist/auth-storage/shared.d.ts.map +1 -1
- package/dist/index.d.ts +2 -1
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +1721 -393
- package/dist/index.js.map +17 -14
- package/dist/migrations.d.ts.map +1 -1
- package/dist/pg-ai-cost-storage.d.ts +2 -1
- package/dist/pg-ai-cost-storage.d.ts.map +1 -1
- package/dist/pg-ai-limit-storage.d.ts +34 -0
- package/dist/pg-ai-limit-storage.d.ts.map +1 -0
- package/dist/pg-ai-usage-storage.d.ts +28 -0
- package/dist/pg-ai-usage-storage.d.ts.map +1 -1
- package/package.json +4 -4
- package/src/agents/rows.ts +38 -9
- package/src/agents/runs.ts +148 -11
- package/src/agents/threads.ts +1 -17
- package/src/ai-cost-groups.ts +102 -0
- package/src/auth-storage/auth-storage.ts +84 -0
- package/src/auth-storage/device-authorizations.ts +131 -0
- package/src/auth-storage/rows.ts +41 -0
- package/src/auth-storage/shared.ts +1 -0
- package/src/index.ts +8 -1
- package/src/migrations/029-model-accounting.sql +24 -0
- package/src/migrations/030-ai-usage-limits.sql +94 -0
- package/src/migrations/031-subagent-runs.sql +97 -0
- package/src/migrations/032-conversation-run-spec.sql +23 -0
- package/src/migrations/033-retire-agent-definitions.sql +28 -0
- package/src/migrations/034-device-authorizations.sql +24 -0
- package/src/migrations.ts +18 -0
- package/src/pg-ai-cost-storage.ts +71 -24
- package/src/pg-ai-limit-storage.ts +901 -0
- package/src/pg-ai-usage-storage.ts +34 -0
- package/src/pg-workflow-run-storage.ts +8 -8
package/dist/index.js
CHANGED
|
@@ -104,7 +104,6 @@ function rowToThreadRecord(row) {
|
|
|
104
104
|
return {
|
|
105
105
|
id: row.id,
|
|
106
106
|
projectId: row.project_id,
|
|
107
|
-
agentId: row.agent_id,
|
|
108
107
|
ownerPrincipal: { type: row.owner_principal_type, id: row.owner_principal_id },
|
|
109
108
|
title: row.title ?? undefined,
|
|
110
109
|
status: row.status,
|
|
@@ -116,13 +115,10 @@ function rowToThreadRecord(row) {
|
|
|
116
115
|
};
|
|
117
116
|
}
|
|
118
117
|
function rowToRunRecord(row) {
|
|
119
|
-
|
|
118
|
+
const base = {
|
|
120
119
|
id: row.id,
|
|
121
120
|
projectId: row.project_id,
|
|
122
121
|
executionId: row.execution_id,
|
|
123
|
-
threadId: row.thread_id,
|
|
124
|
-
agentId: row.agent_id,
|
|
125
|
-
triggerMessageId: row.trigger_message_id,
|
|
126
122
|
requesterGroupIds: typeof row.requester_group_ids === "string" ? JSON.parse(row.requester_group_ids) : row.requester_group_ids,
|
|
127
123
|
status: row.status,
|
|
128
124
|
modelId: row.model_id ?? undefined,
|
|
@@ -138,6 +134,29 @@ function rowToRunRecord(row) {
|
|
|
138
134
|
startedAt: row.started_at ? new Date(row.started_at) : undefined,
|
|
139
135
|
completedAt: row.completed_at ? new Date(row.completed_at) : undefined
|
|
140
136
|
};
|
|
137
|
+
if (row.kind === "conversation") {
|
|
138
|
+
if (!row.thread_id || !row.trigger_message_id) {
|
|
139
|
+
throw new Error(`[SixbPg] Conversational Agent run '${row.id}' has invalid persisted fields.`);
|
|
140
|
+
}
|
|
141
|
+
return {
|
|
142
|
+
...base,
|
|
143
|
+
kind: "conversation",
|
|
144
|
+
threadId: row.thread_id,
|
|
145
|
+
triggerMessageId: row.trigger_message_id,
|
|
146
|
+
...row.spec === null ? {} : { spec: parseJsonColumn(row.spec) }
|
|
147
|
+
};
|
|
148
|
+
}
|
|
149
|
+
if (!row.parent_run_id || !row.spawn_key || row.spec === null) {
|
|
150
|
+
throw new Error(`[SixbPg] Subagent run '${row.id}' has invalid persisted fields.`);
|
|
151
|
+
}
|
|
152
|
+
return {
|
|
153
|
+
...base,
|
|
154
|
+
kind: "subagent",
|
|
155
|
+
parentRunId: row.parent_run_id,
|
|
156
|
+
spawnKey: row.spawn_key,
|
|
157
|
+
spec: parseJsonColumn(row.spec),
|
|
158
|
+
...row.result === null ? {} : { result: parseJsonColumn(row.result) }
|
|
159
|
+
};
|
|
141
160
|
}
|
|
142
161
|
function rowToMessageRecord(row) {
|
|
143
162
|
return {
|
|
@@ -193,6 +212,9 @@ function normalizeDiagnostics(value) {
|
|
|
193
212
|
}
|
|
194
213
|
return typeof value === "string" ? JSON.parse(value) : value;
|
|
195
214
|
}
|
|
215
|
+
function parseJsonColumn(value) {
|
|
216
|
+
return typeof value === "string" ? JSON.parse(value) : value;
|
|
217
|
+
}
|
|
196
218
|
async function queryAgentList(input) {
|
|
197
219
|
const where = `WHERE ${input.whereClauses.join(" AND ")}`;
|
|
198
220
|
const offset = input.offset ?? 0;
|
|
@@ -523,7 +545,13 @@ class PgAgentMessageStore {
|
|
|
523
545
|
}
|
|
524
546
|
|
|
525
547
|
// src/agents/runs.ts
|
|
526
|
-
import {
|
|
548
|
+
import {
|
|
549
|
+
assertAgentRunExecution,
|
|
550
|
+
assertConversationAgentRunSpec,
|
|
551
|
+
assertCreateSubagentRunInput,
|
|
552
|
+
assertSubagentRunResult,
|
|
553
|
+
subagentRunMatchesCreateInput
|
|
554
|
+
} from "@sixb/core/internal/agent-run-storage-provider";
|
|
527
555
|
import { normalizeRequesterGroupIds } from "@sixb/core/internal/auth";
|
|
528
556
|
import { serializeSixbFailure } from "@sixb/core/internal/errors";
|
|
529
557
|
import {
|
|
@@ -597,13 +625,13 @@ class PgAgentRunStore {
|
|
|
597
625
|
this.executions = executions;
|
|
598
626
|
}
|
|
599
627
|
async create(input) {
|
|
628
|
+
assertConversationAgentRunSpec(input.spec, "SixbPg");
|
|
600
629
|
const createdAt = input.createdAt ?? new Date;
|
|
601
630
|
await assertAgentRunExecution({
|
|
602
631
|
executions: this.executions,
|
|
603
632
|
projectId: input.projectId,
|
|
604
633
|
executionId: input.executionId,
|
|
605
|
-
runId: input.id
|
|
606
|
-
agentId: input.agentId
|
|
634
|
+
runId: input.id
|
|
607
635
|
});
|
|
608
636
|
try {
|
|
609
637
|
return await runPgTransaction(this.sql, async (tx) => {
|
|
@@ -623,9 +651,10 @@ class PgAgentRunStore {
|
|
|
623
651
|
project_id,
|
|
624
652
|
id,
|
|
625
653
|
execution_id,
|
|
654
|
+
kind,
|
|
626
655
|
thread_id,
|
|
627
|
-
agent_id,
|
|
628
656
|
trigger_message_id,
|
|
657
|
+
spec,
|
|
629
658
|
requester_group_ids,
|
|
630
659
|
status,
|
|
631
660
|
attempt,
|
|
@@ -634,9 +663,10 @@ class PgAgentRunStore {
|
|
|
634
663
|
${input.projectId},
|
|
635
664
|
${input.id},
|
|
636
665
|
${input.executionId},
|
|
666
|
+
${"conversation"},
|
|
637
667
|
${input.threadId},
|
|
638
|
-
${input.agentId},
|
|
639
668
|
${input.triggerMessageId},
|
|
669
|
+
${JSON.stringify(input.spec)}::text::jsonb,
|
|
640
670
|
${JSON.stringify(normalizeRequesterGroupIds(input.requesterGroupIds))}::text::jsonb,
|
|
641
671
|
${"queued"},
|
|
642
672
|
${0},
|
|
@@ -648,7 +678,11 @@ class PgAgentRunStore {
|
|
|
648
678
|
UPDATE agent_threads SET active_run_id = ${input.id}, updated_at = ${createdAt}
|
|
649
679
|
WHERE project_id = ${input.projectId} AND id = ${input.threadId}
|
|
650
680
|
`;
|
|
651
|
-
|
|
681
|
+
const created = rowToRunRecord(row);
|
|
682
|
+
if (created.kind !== "conversation") {
|
|
683
|
+
throw new Error(`[SixbPg] Agent run insert '${input.id}' returned the wrong run kind.`);
|
|
684
|
+
}
|
|
685
|
+
return created;
|
|
652
686
|
});
|
|
653
687
|
} catch (error) {
|
|
654
688
|
if (isUniqueViolation(error)) {
|
|
@@ -657,6 +691,80 @@ class PgAgentRunStore {
|
|
|
657
691
|
throw error;
|
|
658
692
|
}
|
|
659
693
|
}
|
|
694
|
+
async createSubagent(input) {
|
|
695
|
+
assertCreateSubagentRunInput(input, "SixbPg");
|
|
696
|
+
const execution = await assertAgentRunExecution({
|
|
697
|
+
executions: this.executions,
|
|
698
|
+
projectId: input.projectId,
|
|
699
|
+
executionId: input.executionId,
|
|
700
|
+
runId: input.id
|
|
701
|
+
});
|
|
702
|
+
try {
|
|
703
|
+
return await runPgTransaction(this.sql, async (tx) => {
|
|
704
|
+
const [parent] = await tx`
|
|
705
|
+
SELECT * FROM agent_runs
|
|
706
|
+
WHERE project_id = ${input.projectId} AND id = ${input.parentRunId}
|
|
707
|
+
FOR UPDATE
|
|
708
|
+
`;
|
|
709
|
+
if (!parent || parent.kind !== "conversation") {
|
|
710
|
+
throw new AgentStorageError3("run_not_found", `[SixbPg] Parent Agent run '${input.parentRunId}' was not found.`);
|
|
711
|
+
}
|
|
712
|
+
const [existingRow] = await tx`
|
|
713
|
+
SELECT * FROM agent_runs
|
|
714
|
+
WHERE project_id = ${input.projectId} AND id = ${input.id}
|
|
715
|
+
`;
|
|
716
|
+
if (existingRow) {
|
|
717
|
+
const existing = rowToRunRecord(existingRow);
|
|
718
|
+
if (subagentRunMatchesCreateInput(existing, input))
|
|
719
|
+
return existing;
|
|
720
|
+
throw new AgentStorageError3("duplicate_id", `[SixbPg] Subagent run '${input.id}' already exists with different immutable inputs.`);
|
|
721
|
+
}
|
|
722
|
+
if (parent.status !== "running") {
|
|
723
|
+
throw new AgentStorageError3("invalid_state", `[SixbPg] Parent Agent run '${parent.id}' is not running (status '${parent.status}').`);
|
|
724
|
+
}
|
|
725
|
+
if (execution.source.type !== "execution" || execution.source.executionId !== parent.execution_id) {
|
|
726
|
+
throw new AgentStorageError3("invalid_input", `[SixbPg] Subagent execution '${execution.id}' is not a child of parent execution '${parent.execution_id}'.`);
|
|
727
|
+
}
|
|
728
|
+
if (parent.execution_token !== input.parentExecutionToken) {
|
|
729
|
+
throw new AgentStorageError3("execution_lost", `[SixbPg] Parent Agent run '${parent.id}' is no longer owned by this execution.`);
|
|
730
|
+
}
|
|
731
|
+
const parentRecord = rowToRunRecord(parent);
|
|
732
|
+
const [countRow] = await tx`
|
|
733
|
+
SELECT COUNT(*)::bigint AS count FROM agent_runs
|
|
734
|
+
WHERE project_id = ${input.projectId}
|
|
735
|
+
AND parent_run_id = ${input.parentRunId}
|
|
736
|
+
AND status IN (${"queued"}, ${"running"})
|
|
737
|
+
`;
|
|
738
|
+
const activeChildren = Number(countRow?.count ?? 0);
|
|
739
|
+
if (activeChildren >= input.maxActiveChildren) {
|
|
740
|
+
throw new AgentStorageError3("active_child_limit", `[SixbPg] Agent run '${parent.id}' already has ${activeChildren} active subagents.`);
|
|
741
|
+
}
|
|
742
|
+
const [row] = await tx`
|
|
743
|
+
INSERT INTO agent_runs (
|
|
744
|
+
project_id, id, execution_id, kind, parent_run_id, spawn_key, spec,
|
|
745
|
+
requester_group_ids, status, attempt, created_at
|
|
746
|
+
) VALUES (
|
|
747
|
+
${input.projectId}, ${input.id}, ${input.executionId}, ${"subagent"},
|
|
748
|
+
${input.parentRunId}, ${input.spawnKey},
|
|
749
|
+
${JSON.stringify(input.spec)}::text::jsonb,
|
|
750
|
+
${JSON.stringify(parentRecord.requesterGroupIds)}::text::jsonb,
|
|
751
|
+
${"queued"}, ${0}, ${input.createdAt ?? new Date}
|
|
752
|
+
)
|
|
753
|
+
RETURNING *
|
|
754
|
+
`;
|
|
755
|
+
const created = rowToRunRecord(row);
|
|
756
|
+
if (created.kind !== "subagent") {
|
|
757
|
+
throw new Error(`[SixbPg] Subagent insert '${input.id}' returned the wrong run kind.`);
|
|
758
|
+
}
|
|
759
|
+
return created;
|
|
760
|
+
});
|
|
761
|
+
} catch (error) {
|
|
762
|
+
if (isUniqueViolation(error)) {
|
|
763
|
+
throw new AgentStorageError3("duplicate_id", `[SixbPg] Subagent run '${input.id}' conflicts with an existing durable spawn.`);
|
|
764
|
+
}
|
|
765
|
+
throw error;
|
|
766
|
+
}
|
|
767
|
+
}
|
|
660
768
|
async start(input) {
|
|
661
769
|
const startedAt = input.startedAt ?? new Date;
|
|
662
770
|
return runPgTransaction(this.sql, async (tx) => {
|
|
@@ -731,6 +839,12 @@ class PgAgentRunStore {
|
|
|
731
839
|
const errorValue = input.status === "succeeded" || input.error === undefined ? null : serializeSixbFailure(input.error, AGENT_RUN_FAILURE_CODES2);
|
|
732
840
|
return runPgTransaction(this.sql, async (tx) => {
|
|
733
841
|
const run = await this.lockRunning(tx, input.projectId, input.id);
|
|
842
|
+
if (input.status === "succeeded" && run.kind === "subagent") {
|
|
843
|
+
assertSubagentRunResult(input.result, run.id, "SixbPg");
|
|
844
|
+
}
|
|
845
|
+
if (input.status === "succeeded" && run.kind === "conversation" && input.result !== undefined) {
|
|
846
|
+
throw new AgentStorageError3("invalid_input", `[SixbPg] Conversational Agent run '${run.id}' cannot persist a subagent result.`);
|
|
847
|
+
}
|
|
734
848
|
if (run.execution_token !== input.executionToken) {
|
|
735
849
|
throw new AgentStorageError3("execution_lost", `[SixbPg] Execution token is no longer current on agent run '${input.id}'.`);
|
|
736
850
|
}
|
|
@@ -742,6 +856,7 @@ class PgAgentRunStore {
|
|
|
742
856
|
finish_reason = ${input.finishReason ?? null},
|
|
743
857
|
error = ${errorValue}::text::jsonb,
|
|
744
858
|
diagnostics = ${input.diagnostics === undefined ? null : JSON.stringify(input.diagnostics)}::text::jsonb,
|
|
859
|
+
result = ${input.status === "succeeded" && input.result !== undefined ? JSON.stringify(input.result) : null}::text::jsonb,
|
|
745
860
|
execution_token = ${null},
|
|
746
861
|
execution_queue_lease_expires_at = ${null},
|
|
747
862
|
completed_at = ${completedAt}
|
|
@@ -776,19 +891,23 @@ class PgAgentRunStore {
|
|
|
776
891
|
});
|
|
777
892
|
}
|
|
778
893
|
async list(input) {
|
|
779
|
-
if (hasEmptyStatuses(input)) {
|
|
894
|
+
if (hasEmptyStatuses(input) || input.kinds?.length === 0) {
|
|
780
895
|
return { runs: [], hasMore: false, total: 0 };
|
|
781
896
|
}
|
|
782
897
|
const whereClauses = ["project_id = $1"];
|
|
783
898
|
const params = [input.projectId];
|
|
784
899
|
let index = 2;
|
|
900
|
+
if (input.kinds) {
|
|
901
|
+
whereClauses.push(`kind = ANY($${index++}::text[])`);
|
|
902
|
+
params.push(input.kinds);
|
|
903
|
+
}
|
|
785
904
|
if (input.threadId) {
|
|
786
905
|
whereClauses.push(`thread_id = $${index++}`);
|
|
787
906
|
params.push(input.threadId);
|
|
788
907
|
}
|
|
789
|
-
if (input.
|
|
790
|
-
whereClauses.push(`
|
|
791
|
-
params.push(input.
|
|
908
|
+
if (input.parentRunId) {
|
|
909
|
+
whereClauses.push(`parent_run_id = $${index++}`);
|
|
910
|
+
params.push(input.parentRunId);
|
|
792
911
|
}
|
|
793
912
|
index = appendRunListFilters(whereClauses, params, index, input, "COALESCE(started_at, created_at)");
|
|
794
913
|
const { rows, total, hasMore } = await queryRunList({
|
|
@@ -821,6 +940,8 @@ class PgAgentRunStore {
|
|
|
821
940
|
return row;
|
|
822
941
|
}
|
|
823
942
|
async releaseThread(sql, run, completedAt) {
|
|
943
|
+
if (run.kind !== "conversation")
|
|
944
|
+
return;
|
|
824
945
|
await sql`
|
|
825
946
|
UPDATE agent_threads SET active_run_id = ${null}, updated_at = ${completedAt}
|
|
826
947
|
WHERE project_id = ${run.project_id} AND id = ${run.thread_id} AND active_run_id = ${run.id}
|
|
@@ -845,7 +966,6 @@ class PgAgentThreadStore {
|
|
|
845
966
|
INSERT INTO agent_threads (
|
|
846
967
|
project_id,
|
|
847
968
|
id,
|
|
848
|
-
agent_id,
|
|
849
969
|
owner_principal_type,
|
|
850
970
|
owner_principal_id,
|
|
851
971
|
title,
|
|
@@ -855,7 +975,6 @@ class PgAgentThreadStore {
|
|
|
855
975
|
) VALUES (
|
|
856
976
|
${input.projectId},
|
|
857
977
|
${input.id},
|
|
858
|
-
${input.agentId},
|
|
859
978
|
${input.ownerPrincipal.type},
|
|
860
979
|
${input.ownerPrincipal.id},
|
|
861
980
|
${input.title ?? null},
|
|
@@ -880,21 +999,12 @@ class PgAgentThreadStore {
|
|
|
880
999
|
return row ? rowToThreadRecord(row) : null;
|
|
881
1000
|
}
|
|
882
1001
|
async list(input) {
|
|
883
|
-
if (input.statuses !== undefined && input.statuses.length === 0
|
|
1002
|
+
if (input.statuses !== undefined && input.statuses.length === 0) {
|
|
884
1003
|
return { threads: [], hasMore: false, total: 0 };
|
|
885
1004
|
}
|
|
886
1005
|
const whereClauses = ["project_id = $1"];
|
|
887
1006
|
const params = [input.projectId];
|
|
888
1007
|
let index = 2;
|
|
889
|
-
if (input.agentId) {
|
|
890
|
-
whereClauses.push(`agent_id = $${index++}`);
|
|
891
|
-
params.push(input.agentId);
|
|
892
|
-
}
|
|
893
|
-
if (input.agentIds) {
|
|
894
|
-
const placeholders = input.agentIds.map(() => `$${index++}`);
|
|
895
|
-
whereClauses.push(`agent_id IN (${placeholders.join(", ")})`);
|
|
896
|
-
params.push(...input.agentIds);
|
|
897
|
-
}
|
|
898
1008
|
if (input.statuses) {
|
|
899
1009
|
const placeholders = input.statuses.map(() => `$${index++}`);
|
|
900
1010
|
whereClauses.push(`status IN (${placeholders.join(", ")})`);
|
|
@@ -933,7 +1043,7 @@ class PgAgentStorage {
|
|
|
933
1043
|
}
|
|
934
1044
|
}
|
|
935
1045
|
// src/auth-storage/auth-storage.ts
|
|
936
|
-
import { AuthStorageError as
|
|
1046
|
+
import { AuthStorageError as AuthStorageError11 } from "@sixb/core/storage";
|
|
937
1047
|
|
|
938
1048
|
// src/auth-storage/access-tokens.ts
|
|
939
1049
|
import { AuthStorageError as AuthStorageError2 } from "@sixb/core/storage";
|
|
@@ -1076,6 +1186,25 @@ function rowToOidcAuthorizationAttemptRecord(row) {
|
|
|
1076
1186
|
consumedAt: row.consumed_at ? toDate(row.consumed_at) : undefined
|
|
1077
1187
|
};
|
|
1078
1188
|
}
|
|
1189
|
+
function rowToDeviceAuthorizationRecord(row) {
|
|
1190
|
+
return {
|
|
1191
|
+
id: row.id,
|
|
1192
|
+
projectId: row.project_id,
|
|
1193
|
+
deviceCodeHash: row.device_code_hash,
|
|
1194
|
+
userCode: row.user_code,
|
|
1195
|
+
clientName: row.client_name,
|
|
1196
|
+
tokenName: row.token_name,
|
|
1197
|
+
tokenExpiresAt: toDate(row.token_expires_at),
|
|
1198
|
+
status: row.status,
|
|
1199
|
+
approvedUserId: row.approved_user_id ?? undefined,
|
|
1200
|
+
approvedSessionId: row.approved_session_id ?? undefined,
|
|
1201
|
+
createdAt: toDate(row.created_at),
|
|
1202
|
+
expiresAt: toDate(row.expires_at),
|
|
1203
|
+
approvedAt: row.approved_at ? toDate(row.approved_at) : undefined,
|
|
1204
|
+
deniedAt: row.denied_at ? toDate(row.denied_at) : undefined,
|
|
1205
|
+
consumedAt: row.consumed_at ? toDate(row.consumed_at) : undefined
|
|
1206
|
+
};
|
|
1207
|
+
}
|
|
1079
1208
|
function serializeOptionalRecord(value) {
|
|
1080
1209
|
return value ? JSON.stringify(value) : null;
|
|
1081
1210
|
}
|
|
@@ -1742,6 +1871,102 @@ function assertKindMatchesSubject(input) {
|
|
|
1742
1871
|
throw new AuthStorageError2("invalid_input", `[Sixb] Access token kind '${input.kind}' cannot target subject type '${input.subjectType}'.`);
|
|
1743
1872
|
}
|
|
1744
1873
|
|
|
1874
|
+
// src/auth-storage/device-authorizations.ts
|
|
1875
|
+
import { AuthStorageError as AuthStorageError3, MAX_PENDING_DEVICE_AUTHORIZATIONS } from "@sixb/core/storage";
|
|
1876
|
+
class PgAuthDeviceAuthorizationStore {
|
|
1877
|
+
sql;
|
|
1878
|
+
constructor(sql) {
|
|
1879
|
+
this.sql = sql;
|
|
1880
|
+
}
|
|
1881
|
+
async create(input) {
|
|
1882
|
+
return runPgTransaction(this.sql, async (tx) => {
|
|
1883
|
+
await lockAdvisoryKeys(tx, [authLockKey("device-authorizations", input.projectId)]);
|
|
1884
|
+
await tx`
|
|
1885
|
+
DELETE FROM auth_device_authorizations
|
|
1886
|
+
WHERE project_id = ${input.projectId} AND expires_at <= ${input.createdAt}
|
|
1887
|
+
`;
|
|
1888
|
+
const [pending] = await tx`
|
|
1889
|
+
SELECT COUNT(*)::bigint AS count FROM auth_device_authorizations
|
|
1890
|
+
WHERE project_id = ${input.projectId} AND status = 'pending'
|
|
1891
|
+
`;
|
|
1892
|
+
if (Number(pending?.count ?? 0) >= MAX_PENDING_DEVICE_AUTHORIZATIONS) {
|
|
1893
|
+
throw new AuthStorageError3("device_authorization_limit_exceeded", `[Sixb] Project '${input.projectId}' has too many pending device authorizations.`);
|
|
1894
|
+
}
|
|
1895
|
+
try {
|
|
1896
|
+
const [row] = await tx`
|
|
1897
|
+
INSERT INTO auth_device_authorizations (
|
|
1898
|
+
project_id, id, device_code_hash, user_code, client_name, token_name,
|
|
1899
|
+
token_expires_at, status, created_at, expires_at
|
|
1900
|
+
) VALUES (
|
|
1901
|
+
${assertNonEmpty(input.projectId, "Project id")},
|
|
1902
|
+
${assertNonEmpty(input.id, "Device authorization id")},
|
|
1903
|
+
${assertNonEmpty(input.deviceCodeHash, "Device code hash")},
|
|
1904
|
+
${assertNonEmpty(input.userCode, "User code")},
|
|
1905
|
+
${assertNonEmpty(input.clientName, "Client name")},
|
|
1906
|
+
${assertNonEmpty(input.tokenName, "Token name")},
|
|
1907
|
+
${input.tokenExpiresAt}, 'pending', ${input.createdAt}, ${input.expiresAt}
|
|
1908
|
+
) RETURNING *
|
|
1909
|
+
`;
|
|
1910
|
+
return rowToDeviceAuthorizationRecord(row);
|
|
1911
|
+
} catch (error) {
|
|
1912
|
+
mapUniqueConstraintError(error, "duplicate_device_authorization", `[Sixb] Device authorization '${input.id}' already exists for project '${input.projectId}'.`);
|
|
1913
|
+
}
|
|
1914
|
+
});
|
|
1915
|
+
}
|
|
1916
|
+
async getById(params) {
|
|
1917
|
+
return getPgDeviceAuthorizationById(this.sql, params.projectId, params.id);
|
|
1918
|
+
}
|
|
1919
|
+
async getByUserCode(params) {
|
|
1920
|
+
const [row] = await this.sql`
|
|
1921
|
+
SELECT * FROM auth_device_authorizations
|
|
1922
|
+
WHERE project_id = ${params.projectId} AND user_code = ${params.userCode}
|
|
1923
|
+
`;
|
|
1924
|
+
return row ? rowToDeviceAuthorizationRecord(row) : null;
|
|
1925
|
+
}
|
|
1926
|
+
async approve(params) {
|
|
1927
|
+
const [row] = await this.sql`
|
|
1928
|
+
UPDATE auth_device_authorizations
|
|
1929
|
+
SET status = 'approved', approved_user_id = ${params.userId},
|
|
1930
|
+
approved_session_id = ${params.sessionId}, approved_at = ${params.approvedAt}
|
|
1931
|
+
WHERE project_id = ${params.projectId} AND id = ${params.id}
|
|
1932
|
+
AND status = 'pending' AND expires_at > ${params.approvedAt}
|
|
1933
|
+
AND EXISTS (
|
|
1934
|
+
SELECT 1 FROM auth_sessions
|
|
1935
|
+
WHERE project_id = ${params.projectId} AND id = ${params.sessionId}
|
|
1936
|
+
AND user_id = ${params.userId} AND revoked_at IS NULL
|
|
1937
|
+
AND expires_at > ${params.approvedAt}
|
|
1938
|
+
AND (absolute_expires_at IS NULL OR absolute_expires_at > ${params.approvedAt})
|
|
1939
|
+
)
|
|
1940
|
+
RETURNING *
|
|
1941
|
+
`;
|
|
1942
|
+
if (!row)
|
|
1943
|
+
throw invalidTransition(params.id, params.projectId);
|
|
1944
|
+
return rowToDeviceAuthorizationRecord(row);
|
|
1945
|
+
}
|
|
1946
|
+
async deny(params) {
|
|
1947
|
+
const [row] = await this.sql`
|
|
1948
|
+
UPDATE auth_device_authorizations
|
|
1949
|
+
SET status = 'denied', denied_at = ${params.deniedAt}
|
|
1950
|
+
WHERE project_id = ${params.projectId} AND id = ${params.id}
|
|
1951
|
+
AND status = 'pending' AND expires_at > ${params.deniedAt}
|
|
1952
|
+
RETURNING *
|
|
1953
|
+
`;
|
|
1954
|
+
if (!row)
|
|
1955
|
+
throw invalidTransition(params.id, params.projectId);
|
|
1956
|
+
return rowToDeviceAuthorizationRecord(row);
|
|
1957
|
+
}
|
|
1958
|
+
}
|
|
1959
|
+
async function getPgDeviceAuthorizationById(sql, projectId, id) {
|
|
1960
|
+
const [row] = await sql`
|
|
1961
|
+
SELECT * FROM auth_device_authorizations
|
|
1962
|
+
WHERE project_id = ${projectId} AND id = ${id}
|
|
1963
|
+
`;
|
|
1964
|
+
return row ? rowToDeviceAuthorizationRecord(row) : null;
|
|
1965
|
+
}
|
|
1966
|
+
function invalidTransition(id, projectId) {
|
|
1967
|
+
return new AuthStorageError3("invalid_device_authorization", `[Sixb] Device authorization '${id}' is no longer pending for project '${projectId}'.`);
|
|
1968
|
+
}
|
|
1969
|
+
|
|
1745
1970
|
// src/auth-storage/group-memberships.ts
|
|
1746
1971
|
class PgAuthGroupMembershipStore {
|
|
1747
1972
|
sql;
|
|
@@ -1832,7 +2057,7 @@ class PgAuthUserIdentityStore {
|
|
|
1832
2057
|
}
|
|
1833
2058
|
|
|
1834
2059
|
// src/auth-storage/invitations.ts
|
|
1835
|
-
import { AuthStorageError as
|
|
2060
|
+
import { AuthStorageError as AuthStorageError4 } from "@sixb/core/storage";
|
|
1836
2061
|
class PgAuthInvitationStore {
|
|
1837
2062
|
sql;
|
|
1838
2063
|
constructor(sql) {
|
|
@@ -1873,7 +2098,7 @@ class PgAuthInvitationStore {
|
|
|
1873
2098
|
return requireInvitationById(tx, { projectId, id: existingActive.id });
|
|
1874
2099
|
}
|
|
1875
2100
|
if (await getInvitationById(tx, { projectId, id })) {
|
|
1876
|
-
throw new
|
|
2101
|
+
throw new AuthStorageError4("duplicate_invitation", `[Sixb] Invitation '${id}' already exists but is not active for project '${projectId}'.`);
|
|
1877
2102
|
}
|
|
1878
2103
|
const createdAt = dateOrNow(input.createdAt);
|
|
1879
2104
|
try {
|
|
@@ -2011,7 +2236,7 @@ class PgAuthInvitationStore {
|
|
|
2011
2236
|
|
|
2012
2237
|
// src/auth-storage/magic-links.ts
|
|
2013
2238
|
import { resolveAuthSessionAudience } from "@sixb/core";
|
|
2014
|
-
import { AuthStorageError as
|
|
2239
|
+
import { AuthStorageError as AuthStorageError5 } from "@sixb/core/storage";
|
|
2015
2240
|
class PgAuthMagicLinkStore {
|
|
2016
2241
|
sql;
|
|
2017
2242
|
constructor(sql) {
|
|
@@ -2027,7 +2252,7 @@ class PgAuthMagicLinkStore {
|
|
|
2027
2252
|
const tokenHash = assertNonEmpty(input.tokenHash, "Magic link token hash");
|
|
2028
2253
|
await lockAdvisoryKeys(tx, [authLockKey("magic-links", projectId, email)]);
|
|
2029
2254
|
if (await getMagicLinkRowById(tx, { projectId, id })) {
|
|
2030
|
-
throw new
|
|
2255
|
+
throw new AuthStorageError5("duplicate_magic_link", `[Sixb] Magic link '${id}' already exists for project '${projectId}'.`);
|
|
2031
2256
|
}
|
|
2032
2257
|
await revokeActiveMagicLinksForEmail(tx, {
|
|
2033
2258
|
projectId,
|
|
@@ -2099,7 +2324,7 @@ class PgAuthMagicLinkStore {
|
|
|
2099
2324
|
|
|
2100
2325
|
// src/auth-storage/oidc-attempts.ts
|
|
2101
2326
|
import { resolveAuthSessionAudience as resolveAuthSessionAudience2 } from "@sixb/core";
|
|
2102
|
-
import { AuthStorageError as
|
|
2327
|
+
import { AuthStorageError as AuthStorageError6 } from "@sixb/core/storage";
|
|
2103
2328
|
class PgAuthOidcAuthorizationAttemptStore {
|
|
2104
2329
|
sql;
|
|
2105
2330
|
constructor(sql) {
|
|
@@ -2114,7 +2339,7 @@ class PgAuthOidcAuthorizationAttemptStore {
|
|
|
2114
2339
|
const nonceHash = assertNonEmpty(input.nonceHash, "OIDC nonce hash");
|
|
2115
2340
|
const codeVerifier = assertNonEmpty(input.codeVerifier, "OIDC code verifier");
|
|
2116
2341
|
if (await getOidcAttemptRowById(this.sql, { projectId, id })) {
|
|
2117
|
-
throw new
|
|
2342
|
+
throw new AuthStorageError6("duplicate_oidc_attempt", `[Sixb] OIDC authorization attempt '${id}' already exists for project '${projectId}'.`);
|
|
2118
2343
|
}
|
|
2119
2344
|
try {
|
|
2120
2345
|
const [row] = await this.sql`
|
|
@@ -2157,7 +2382,7 @@ class PgAuthOidcAuthorizationAttemptStore {
|
|
|
2157
2382
|
}
|
|
2158
2383
|
|
|
2159
2384
|
// src/auth-storage/service-account-group-memberships.ts
|
|
2160
|
-
import { AuthStorageError as
|
|
2385
|
+
import { AuthStorageError as AuthStorageError7 } from "@sixb/core/storage";
|
|
2161
2386
|
class PgAuthServiceAccountGroupMembershipStore {
|
|
2162
2387
|
sql;
|
|
2163
2388
|
constructor(sql) {
|
|
@@ -2174,7 +2399,7 @@ class PgAuthServiceAccountGroupMembershipStore {
|
|
|
2174
2399
|
AND id = ${serviceAccountId}
|
|
2175
2400
|
`;
|
|
2176
2401
|
if (!serviceAccount) {
|
|
2177
|
-
throw new
|
|
2402
|
+
throw new AuthStorageError7("missing_service_account", `[Sixb] Service account '${serviceAccountId}' not found for project '${projectId}'.`);
|
|
2178
2403
|
}
|
|
2179
2404
|
const [row] = await this.sql`
|
|
2180
2405
|
INSERT INTO auth_service_account_group_memberships (
|
|
@@ -2209,7 +2434,7 @@ class PgAuthServiceAccountGroupMembershipStore {
|
|
|
2209
2434
|
AND id = ${serviceAccountId}
|
|
2210
2435
|
`;
|
|
2211
2436
|
if (!serviceAccount) {
|
|
2212
|
-
throw new
|
|
2437
|
+
throw new AuthStorageError7("missing_service_account", `[Sixb] Service account '${serviceAccountId}' not found for project '${projectId}'.`);
|
|
2213
2438
|
}
|
|
2214
2439
|
if (groupIds.length === 0) {
|
|
2215
2440
|
await this.sql`
|
|
@@ -2272,7 +2497,7 @@ class PgAuthServiceAccountGroupMembershipStore {
|
|
|
2272
2497
|
}
|
|
2273
2498
|
|
|
2274
2499
|
// src/auth-storage/service-accounts.ts
|
|
2275
|
-
import { AuthStorageError as
|
|
2500
|
+
import { AuthStorageError as AuthStorageError8 } from "@sixb/core/storage";
|
|
2276
2501
|
class PgAuthServiceAccountStore {
|
|
2277
2502
|
sql;
|
|
2278
2503
|
constructor(sql) {
|
|
@@ -2343,7 +2568,7 @@ class PgAuthServiceAccountStore {
|
|
|
2343
2568
|
RETURNING *
|
|
2344
2569
|
`;
|
|
2345
2570
|
if (!row) {
|
|
2346
|
-
throw new
|
|
2571
|
+
throw new AuthStorageError8("missing_service_account", `[Sixb] Service account '${id}' not found for project '${projectId}'.`);
|
|
2347
2572
|
}
|
|
2348
2573
|
return rowToServiceAccountRecord(row);
|
|
2349
2574
|
}
|
|
@@ -2379,7 +2604,7 @@ class PgAuthServiceAccountStore {
|
|
|
2379
2604
|
}
|
|
2380
2605
|
|
|
2381
2606
|
// src/auth-storage/sessions.ts
|
|
2382
|
-
import { AuthStorageError as
|
|
2607
|
+
import { AuthStorageError as AuthStorageError9 } from "@sixb/core/storage";
|
|
2383
2608
|
class PgAuthSessionStore {
|
|
2384
2609
|
sql;
|
|
2385
2610
|
constructor(sql) {
|
|
@@ -2452,7 +2677,7 @@ class PgAuthSessionStore {
|
|
|
2452
2677
|
return runPgTransaction(this.sql, async (tx) => {
|
|
2453
2678
|
const existing = await getSessionRowById(tx, params, { forUpdate: true });
|
|
2454
2679
|
if (!existing) {
|
|
2455
|
-
throw new
|
|
2680
|
+
throw new AuthStorageError9("missing_session", `[Sixb] Session '${params.id}' not found for project '${params.projectId}'.`);
|
|
2456
2681
|
}
|
|
2457
2682
|
const [row] = await tx`
|
|
2458
2683
|
UPDATE auth_sessions
|
|
@@ -2490,7 +2715,7 @@ class PgAuthSessionStore {
|
|
|
2490
2715
|
}
|
|
2491
2716
|
|
|
2492
2717
|
// src/auth-storage/users.ts
|
|
2493
|
-
import { AuthStorageError as
|
|
2718
|
+
import { AuthStorageError as AuthStorageError10 } from "@sixb/core/storage";
|
|
2494
2719
|
class PgAuthUserStore {
|
|
2495
2720
|
sql;
|
|
2496
2721
|
constructor(sql) {
|
|
@@ -2501,10 +2726,10 @@ class PgAuthUserStore {
|
|
|
2501
2726
|
const projectId = assertNonEmpty(input.projectId, "Project id");
|
|
2502
2727
|
const email = normalizeEmail(input.email);
|
|
2503
2728
|
if (await getUserRowById(this.sql, { projectId, id })) {
|
|
2504
|
-
throw new
|
|
2729
|
+
throw new AuthStorageError10("duplicate_user", `[Sixb] User '${id}' already exists for project '${projectId}'.`);
|
|
2505
2730
|
}
|
|
2506
2731
|
if (await getUserRowByEmail(this.sql, { projectId, email })) {
|
|
2507
|
-
throw new
|
|
2732
|
+
throw new AuthStorageError10("duplicate_user", `[Sixb] User email '${email}' already exists for project '${projectId}'.`);
|
|
2508
2733
|
}
|
|
2509
2734
|
const createdAt = dateOrNow(input.createdAt);
|
|
2510
2735
|
const updatedAt = input.updatedAt ? new Date(input.updatedAt) : createdAt;
|
|
@@ -2550,7 +2775,7 @@ class PgAuthUserStore {
|
|
|
2550
2775
|
id: input.id
|
|
2551
2776
|
});
|
|
2552
2777
|
if (!existing) {
|
|
2553
|
-
throw new
|
|
2778
|
+
throw new AuthStorageError10("missing_user", `[Sixb] User '${input.id}' not found for project '${input.projectId}'.`);
|
|
2554
2779
|
}
|
|
2555
2780
|
const updatedAt = dateOrNow(input.updatedAt);
|
|
2556
2781
|
const [row] = await this.sql`
|
|
@@ -2570,7 +2795,7 @@ class PgAuthUserStore {
|
|
|
2570
2795
|
id: input.id
|
|
2571
2796
|
});
|
|
2572
2797
|
if (!existing) {
|
|
2573
|
-
throw new
|
|
2798
|
+
throw new AuthStorageError10("missing_user", `[Sixb] User '${input.id}' not found for project '${input.projectId}'.`);
|
|
2574
2799
|
}
|
|
2575
2800
|
const updatedAt = dateOrNow(input.updatedAt);
|
|
2576
2801
|
const [row] = await this.sql`
|
|
@@ -2627,6 +2852,7 @@ class PgAuthStorage {
|
|
|
2627
2852
|
groupMemberships;
|
|
2628
2853
|
magicLinks;
|
|
2629
2854
|
oidcAuthorizationAttempts;
|
|
2855
|
+
deviceAuthorizations;
|
|
2630
2856
|
constructor(options) {
|
|
2631
2857
|
this.sql = options.sql;
|
|
2632
2858
|
this.users = new PgAuthUserStore(this.sql);
|
|
@@ -2639,6 +2865,34 @@ class PgAuthStorage {
|
|
|
2639
2865
|
this.groupMemberships = new PgAuthGroupMembershipStore(this.sql);
|
|
2640
2866
|
this.magicLinks = new PgAuthMagicLinkStore(this.sql);
|
|
2641
2867
|
this.oidcAuthorizationAttempts = new PgAuthOidcAuthorizationAttemptStore(this.sql);
|
|
2868
|
+
this.deviceAuthorizations = new PgAuthDeviceAuthorizationStore(this.sql);
|
|
2869
|
+
}
|
|
2870
|
+
async completeDeviceAuthorization(input) {
|
|
2871
|
+
return runPgTransaction(this.sql, async (tx) => {
|
|
2872
|
+
const authorization = await getPgDeviceAuthorizationById(tx, input.projectId, input.id);
|
|
2873
|
+
if (!authorization) {
|
|
2874
|
+
throw new AuthStorageError11("missing_device_authorization", `[Sixb] Device authorization '${input.id}' not found for project '${input.projectId}'.`);
|
|
2875
|
+
}
|
|
2876
|
+
assertCompletableDeviceAuthorization(authorization, input);
|
|
2877
|
+
const accessToken = await new PgAuthAccessTokenStore(tx).create(input.accessToken);
|
|
2878
|
+
const [consumed] = await tx`
|
|
2879
|
+
UPDATE auth_device_authorizations
|
|
2880
|
+
SET status = 'consumed', consumed_at = ${input.completedAt}
|
|
2881
|
+
WHERE project_id = ${input.projectId} AND id = ${input.id} AND status = 'approved'
|
|
2882
|
+
RETURNING *
|
|
2883
|
+
`;
|
|
2884
|
+
if (!consumed) {
|
|
2885
|
+
throw new AuthStorageError11("invalid_device_authorization", "[Sixb] Device authorization was already consumed.");
|
|
2886
|
+
}
|
|
2887
|
+
return {
|
|
2888
|
+
authorization: {
|
|
2889
|
+
...authorization,
|
|
2890
|
+
status: "consumed",
|
|
2891
|
+
consumedAt: new Date(input.completedAt)
|
|
2892
|
+
},
|
|
2893
|
+
accessToken
|
|
2894
|
+
};
|
|
2895
|
+
});
|
|
2642
2896
|
}
|
|
2643
2897
|
async completeMagicLinkSignIn(input) {
|
|
2644
2898
|
const result = await runPgTransaction(this.sql, async (tx) => {
|
|
@@ -2649,7 +2903,7 @@ class PgAuthStorage {
|
|
|
2649
2903
|
id: input.magicLinkId
|
|
2650
2904
|
});
|
|
2651
2905
|
if (!initialMagicLink) {
|
|
2652
|
-
throw new
|
|
2906
|
+
throw new AuthStorageError11("missing_magic_link", `[Sixb] Magic link '${input.magicLinkId}' not found for project '${projectId}'.`);
|
|
2653
2907
|
}
|
|
2654
2908
|
const email = normalizeEmail(initialMagicLink.email);
|
|
2655
2909
|
await lockAdvisoryKeys(tx, [
|
|
@@ -2682,7 +2936,7 @@ class PgAuthStorage {
|
|
|
2682
2936
|
consumedAt: completedAt
|
|
2683
2937
|
});
|
|
2684
2938
|
return {
|
|
2685
|
-
error: new
|
|
2939
|
+
error: new AuthStorageError11("suspended_user", `[Sixb] User '${existingUserRow.id}' is suspended for project '${projectId}'.`)
|
|
2686
2940
|
};
|
|
2687
2941
|
}
|
|
2688
2942
|
if (shouldCreateUser && !activeInvitation && !input.allowUserCreationWithoutInvitation) {
|
|
@@ -2693,7 +2947,7 @@ class PgAuthStorage {
|
|
|
2693
2947
|
consumedAt: completedAt
|
|
2694
2948
|
});
|
|
2695
2949
|
return {
|
|
2696
|
-
error: new
|
|
2950
|
+
error: new AuthStorageError11("user_creation_not_allowed", `[Sixb] Magic link '${input.magicLinkId}' cannot create a user for project '${projectId}'.`)
|
|
2697
2951
|
};
|
|
2698
2952
|
}
|
|
2699
2953
|
if (shouldCreateUser && !activeInvitation && input.allowUserCreationWithoutInvitation && input.requireNoActiveUsersForUserCreation && await hasActiveUsers(tx, projectId)) {
|
|
@@ -2704,7 +2958,7 @@ class PgAuthStorage {
|
|
|
2704
2958
|
consumedAt: completedAt
|
|
2705
2959
|
});
|
|
2706
2960
|
return {
|
|
2707
|
-
error: new
|
|
2961
|
+
error: new AuthStorageError11("user_creation_not_allowed", `[Sixb] Magic link '${input.magicLinkId}' cannot create a user for project '${projectId}'.`)
|
|
2708
2962
|
};
|
|
2709
2963
|
}
|
|
2710
2964
|
await validateCompleteSessionInput(tx, projectId, input.session);
|
|
@@ -2713,7 +2967,7 @@ class PgAuthStorage {
|
|
|
2713
2967
|
if (shouldCreateUser) {
|
|
2714
2968
|
newUserId = assertNonEmpty(input.newUserId, "User id");
|
|
2715
2969
|
if (await getUserRowById(tx, { projectId, id: newUserId })) {
|
|
2716
|
-
throw new
|
|
2970
|
+
throw new AuthStorageError11("duplicate_user", `[Sixb] User '${newUserId}' already exists for project '${projectId}'.`);
|
|
2717
2971
|
}
|
|
2718
2972
|
}
|
|
2719
2973
|
const sessionUserId = existingUserRow?.id ?? newUserId ?? input.newUserId;
|
|
@@ -2796,37 +3050,37 @@ class PgAuthStorage {
|
|
|
2796
3050
|
if (identity && !userRow) {
|
|
2797
3051
|
await this.consumeOidcAttempt(input, completedAt, projectId, tx);
|
|
2798
3052
|
return {
|
|
2799
|
-
error: new
|
|
3053
|
+
error: new AuthStorageError11("missing_user", `[Sixb] User '${identity.user_id}' not found for linked OIDC identity.`)
|
|
2800
3054
|
};
|
|
2801
3055
|
}
|
|
2802
3056
|
if (!identity && userRow && (!input.autoLinkByVerifiedEmail || !input.emailVerified)) {
|
|
2803
3057
|
await this.consumeOidcAttempt(input, completedAt, projectId, tx);
|
|
2804
3058
|
return {
|
|
2805
|
-
error: new
|
|
3059
|
+
error: new AuthStorageError11("email_link_not_allowed", `[Sixb] OIDC identity cannot auto-link to user '${userRow.id}' for project '${projectId}'.`)
|
|
2806
3060
|
};
|
|
2807
3061
|
}
|
|
2808
3062
|
if (userRow?.status === "suspended") {
|
|
2809
3063
|
await this.consumeOidcAttempt(input, completedAt, projectId, tx);
|
|
2810
3064
|
return {
|
|
2811
|
-
error: new
|
|
3065
|
+
error: new AuthStorageError11("suspended_user", `[Sixb] User '${userRow.id}' is suspended for project '${projectId}'.`)
|
|
2812
3066
|
};
|
|
2813
3067
|
}
|
|
2814
3068
|
if (shouldCreateUser && !input.emailVerified) {
|
|
2815
3069
|
await this.consumeOidcAttempt(input, completedAt, projectId, tx);
|
|
2816
3070
|
return {
|
|
2817
|
-
error: new
|
|
3071
|
+
error: new AuthStorageError11("user_creation_not_allowed", `[Sixb] OIDC authorization attempt '${input.oidcAuthorizationAttemptId}' cannot create a user for project '${projectId}'.`)
|
|
2818
3072
|
};
|
|
2819
3073
|
}
|
|
2820
3074
|
if (shouldCreateUser && !activeInvitation && !input.allowUserCreationWithoutInvitation) {
|
|
2821
3075
|
await this.consumeOidcAttempt(input, completedAt, projectId, tx);
|
|
2822
3076
|
return {
|
|
2823
|
-
error: new
|
|
3077
|
+
error: new AuthStorageError11("user_creation_not_allowed", `[Sixb] OIDC authorization attempt '${input.oidcAuthorizationAttemptId}' cannot create a user for project '${projectId}'.`)
|
|
2824
3078
|
};
|
|
2825
3079
|
}
|
|
2826
3080
|
if (shouldCreateUser && !activeInvitation && input.allowUserCreationWithoutInvitation && input.requireNoActiveUsersForUserCreation && await hasActiveUsers(tx, projectId)) {
|
|
2827
3081
|
await this.consumeOidcAttempt(input, completedAt, projectId, tx);
|
|
2828
3082
|
return {
|
|
2829
|
-
error: new
|
|
3083
|
+
error: new AuthStorageError11("user_creation_not_allowed", `[Sixb] OIDC authorization attempt '${input.oidcAuthorizationAttemptId}' cannot create a user for project '${projectId}'.`)
|
|
2830
3084
|
};
|
|
2831
3085
|
}
|
|
2832
3086
|
await validateCompleteSessionInput(tx, projectId, input.session);
|
|
@@ -2835,7 +3089,7 @@ class PgAuthStorage {
|
|
|
2835
3089
|
if (shouldCreateUser) {
|
|
2836
3090
|
newUserId = assertNonEmpty(input.newUserId, "User id");
|
|
2837
3091
|
if (await getUserRowById(tx, { projectId, id: newUserId })) {
|
|
2838
|
-
throw new
|
|
3092
|
+
throw new AuthStorageError11("duplicate_user", `[Sixb] User '${newUserId}' already exists for project '${projectId}'.`);
|
|
2839
3093
|
}
|
|
2840
3094
|
}
|
|
2841
3095
|
await lockAdvisoryKeys(tx, [
|
|
@@ -2914,7 +3168,7 @@ class PgAuthStorage {
|
|
|
2914
3168
|
id: input.userId
|
|
2915
3169
|
}, { forUpdate: true });
|
|
2916
3170
|
if (!existing) {
|
|
2917
|
-
throw new
|
|
3171
|
+
throw new AuthStorageError11("missing_user", `[Sixb] User '${input.userId}' not found for project '${input.projectId}'.`);
|
|
2918
3172
|
}
|
|
2919
3173
|
const [updated] = await tx`
|
|
2920
3174
|
UPDATE auth_users
|
|
@@ -2934,13 +3188,13 @@ class PgAuthStorage {
|
|
|
2934
3188
|
}
|
|
2935
3189
|
assertMagicLinkUsable(projectId, id, tokenHash, completedAt, row) {
|
|
2936
3190
|
if (!row) {
|
|
2937
|
-
throw new
|
|
3191
|
+
throw new AuthStorageError11("missing_magic_link", `[Sixb] Magic link '${id}' not found for project '${projectId}'.`);
|
|
2938
3192
|
}
|
|
2939
3193
|
assertMagicLinkUsable(projectId, id, tokenHash, completedAt, row);
|
|
2940
3194
|
}
|
|
2941
3195
|
assertOidcAttemptUsable(projectId, id, stateHash, completedAt, row) {
|
|
2942
3196
|
if (!row) {
|
|
2943
|
-
throw new
|
|
3197
|
+
throw new AuthStorageError11("missing_oidc_attempt", `[Sixb] OIDC authorization attempt '${id}' not found for project '${projectId}'.`);
|
|
2944
3198
|
}
|
|
2945
3199
|
assertOidcAttemptUsable(projectId, id, stateHash, completedAt, row);
|
|
2946
3200
|
}
|
|
@@ -3086,7 +3340,21 @@ function assertSignInSessionAudience(projectId, sessionAudience, storedAudience)
|
|
|
3086
3340
|
if (sessionAudience === storedAudience) {
|
|
3087
3341
|
return;
|
|
3088
3342
|
}
|
|
3089
|
-
throw new
|
|
3343
|
+
throw new AuthStorageError11("invalid_input", `[Sixb] Sign-in session audience '${sessionAudience}' does not match stored auth audience '${storedAudience}' for project '${projectId}'.`);
|
|
3344
|
+
}
|
|
3345
|
+
function assertCompletableDeviceAuthorization(authorization, input) {
|
|
3346
|
+
if (authorization.deviceCodeHash !== input.deviceCodeHash) {
|
|
3347
|
+
throw new AuthStorageError11("invalid_device_authorization", "[Sixb] Invalid device authorization.");
|
|
3348
|
+
}
|
|
3349
|
+
if (authorization.status === "pending") {
|
|
3350
|
+
throw new AuthStorageError11("pending_device_authorization", "[Sixb] Device authorization is pending.");
|
|
3351
|
+
}
|
|
3352
|
+
if (authorization.status === "denied") {
|
|
3353
|
+
throw new AuthStorageError11("device_authorization_denied", "[Sixb] Device authorization was denied.");
|
|
3354
|
+
}
|
|
3355
|
+
if (authorization.status !== "approved" || authorization.expiresAt <= input.completedAt || !authorization.approvedUserId || !authorization.approvedSessionId || input.accessToken.projectId !== input.projectId || input.accessToken.subjectType !== "user" || input.accessToken.subjectId !== authorization.approvedUserId || input.accessToken.createdBySessionId !== authorization.approvedSessionId) {
|
|
3356
|
+
throw new AuthStorageError11("invalid_device_authorization", "[Sixb] Device authorization is invalid.");
|
|
3357
|
+
}
|
|
3090
3358
|
}
|
|
3091
3359
|
// src/connector-connection-storage.ts
|
|
3092
3360
|
import {
|
|
@@ -6245,121 +6513,429 @@ ADD COLUMN edited_at JSONB NOT NULL DEFAULT '{}'::jsonb
|
|
|
6245
6513
|
CHECK (jsonb_typeof(edited_at) = 'object');
|
|
6246
6514
|
`;
|
|
6247
6515
|
|
|
6248
|
-
// src/migrations.
|
|
6249
|
-
var
|
|
6250
|
-
|
|
6251
|
-
|
|
6252
|
-
|
|
6253
|
-
|
|
6254
|
-
|
|
6255
|
-
|
|
6256
|
-
|
|
6257
|
-
return [
|
|
6258
|
-
createPostgresMigrator({
|
|
6259
|
-
sql,
|
|
6260
|
-
schemaName,
|
|
6261
|
-
migrations: postgresStorageMigrations
|
|
6262
|
-
})
|
|
6263
|
-
];
|
|
6264
|
-
}
|
|
6265
|
-
function createPostgresMigrator(params) {
|
|
6266
|
-
return {
|
|
6267
|
-
adapterId: params.migrations.adapterId,
|
|
6268
|
-
latestVersion: params.migrations.latestVersion,
|
|
6269
|
-
async status() {
|
|
6270
|
-
return describeMigrationHistory({
|
|
6271
|
-
migrations: params.migrations,
|
|
6272
|
-
rows: await readPostgresHistory(params.sql, params.schemaName, params.migrations.adapterId)
|
|
6273
|
-
});
|
|
6274
|
-
},
|
|
6275
|
-
migrate() {
|
|
6276
|
-
return withPostgresMigrationLock(params, async (sql) => {
|
|
6277
|
-
const session = postgresMigrationSession(sql, params.schemaName);
|
|
6278
|
-
return runMigrationSet({
|
|
6279
|
-
context: session.context,
|
|
6280
|
-
migrations: params.migrations,
|
|
6281
|
-
state: session.state
|
|
6282
|
-
});
|
|
6283
|
-
});
|
|
6284
|
-
}
|
|
6285
|
-
};
|
|
6286
|
-
}
|
|
6287
|
-
async function migratePostgresStorage(sql, schemaName) {
|
|
6288
|
-
await createPostgresMigrator({
|
|
6289
|
-
sql,
|
|
6290
|
-
schemaName,
|
|
6291
|
-
migrations: postgresStorageMigrations
|
|
6292
|
-
}).migrate();
|
|
6293
|
-
}
|
|
6294
|
-
async function dropSchema(sql, schemaName) {
|
|
6295
|
-
await sql.unsafe(`DROP SCHEMA IF EXISTS ${quoteIdent(schemaName)} CASCADE`);
|
|
6296
|
-
}
|
|
6297
|
-
function quoteIdent(identifier) {
|
|
6298
|
-
if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(identifier)) {
|
|
6299
|
-
throw new Error(`[SixbPg] Invalid identifier: ${identifier}`);
|
|
6300
|
-
}
|
|
6301
|
-
return `"${identifier}"`;
|
|
6302
|
-
}
|
|
6303
|
-
function postgresMigrationSession(sql, schemaName) {
|
|
6304
|
-
const schema = quoteIdent(schemaName);
|
|
6305
|
-
const migrationsTable = `${schema}.sixb_migrations`;
|
|
6306
|
-
let active = null;
|
|
6307
|
-
const connection = () => active ?? sql;
|
|
6308
|
-
return {
|
|
6309
|
-
context: {
|
|
6310
|
-
async exec(sqlText) {
|
|
6311
|
-
await connection().unsafe(sqlText);
|
|
6312
|
-
}
|
|
6313
|
-
},
|
|
6314
|
-
state: {
|
|
6315
|
-
async ensure() {
|
|
6316
|
-
await sql.unsafe(`
|
|
6317
|
-
CREATE SCHEMA IF NOT EXISTS ${schema};
|
|
6516
|
+
// src/migrations/029-model-accounting.sql
|
|
6517
|
+
var _029_model_accounting_default = `-- Call-time model configuration and provider identity. Historical identity remains unknown.
|
|
6518
|
+
ALTER TABLE ai_model_call_usage
|
|
6519
|
+
ADD COLUMN requested_reasoning JSONB
|
|
6520
|
+
CHECK (
|
|
6521
|
+
requested_reasoning IS NULL OR
|
|
6522
|
+
jsonb_typeof(requested_reasoning) IN ('string', 'object')
|
|
6523
|
+
);
|
|
6524
|
+
ALTER TABLE ai_model_call_usage ADD COLUMN provider_ids JSONB;
|
|
6318
6525
|
|
|
6319
|
-
|
|
6320
|
-
|
|
6321
|
-
|
|
6322
|
-
|
|
6323
|
-
|
|
6324
|
-
|
|
6325
|
-
|
|
6326
|
-
|
|
6327
|
-
|
|
6328
|
-
|
|
6329
|
-
|
|
6330
|
-
|
|
6331
|
-
|
|
6332
|
-
|
|
6333
|
-
|
|
6334
|
-
|
|
6335
|
-
|
|
6336
|
-
|
|
6337
|
-
|
|
6338
|
-
|
|
6339
|
-
|
|
6340
|
-
|
|
6341
|
-
|
|
6342
|
-
|
|
6343
|
-
|
|
6344
|
-
|
|
6345
|
-
|
|
6346
|
-
|
|
6347
|
-
|
|
6348
|
-
|
|
6349
|
-
|
|
6350
|
-
|
|
6351
|
-
|
|
6352
|
-
|
|
6353
|
-
|
|
6354
|
-
|
|
6355
|
-
|
|
6356
|
-
|
|
6357
|
-
|
|
6358
|
-
|
|
6359
|
-
|
|
6360
|
-
|
|
6361
|
-
|
|
6362
|
-
|
|
6526
|
+
-- Upgrade the shipped valuation reason while retaining the selected-cost projection.
|
|
6527
|
+
ALTER TABLE ai_model_call_valuations
|
|
6528
|
+
DROP CONSTRAINT ai_model_call_valuations_reason_check;
|
|
6529
|
+
|
|
6530
|
+
UPDATE ai_model_call_valuations SET reason = 'missingRateCard'
|
|
6531
|
+
WHERE reason = 'missingCatalogEntry';
|
|
6532
|
+
|
|
6533
|
+
ALTER TABLE ai_model_call_valuations ADD CONSTRAINT ai_model_call_valuations_reason_check
|
|
6534
|
+
CHECK (reason IS NULL OR reason IN (
|
|
6535
|
+
'missingBillingIdentity',
|
|
6536
|
+
'missingRateCard',
|
|
6537
|
+
'missingUsageMeter',
|
|
6538
|
+
'unsupportedPricingDimension',
|
|
6539
|
+
'invalidUsageForFormula'
|
|
6540
|
+
));
|
|
6541
|
+
`;
|
|
6542
|
+
|
|
6543
|
+
// src/migrations/030-ai-usage-limits.sql
|
|
6544
|
+
var _030_ai_usage_limits_default = `CREATE TABLE ai_usage_limits (
|
|
6545
|
+
project_id TEXT NOT NULL CHECK (length(btrim(project_id)) > 0),
|
|
6546
|
+
id TEXT NOT NULL CHECK (length(btrim(id)) > 0),
|
|
6547
|
+
subject_type TEXT NOT NULL CHECK (
|
|
6548
|
+
subject_type IN ('project', 'group', 'user', 'serviceAccount')
|
|
6549
|
+
),
|
|
6550
|
+
subject_id TEXT NOT NULL CHECK (
|
|
6551
|
+
(subject_type = 'project' AND subject_id = '')
|
|
6552
|
+
OR (subject_type <> 'project' AND length(btrim(subject_id)) > 0)
|
|
6553
|
+
),
|
|
6554
|
+
meter TEXT NOT NULL CHECK (meter IN ('tokens.total', 'cost.catalogEstimated')),
|
|
6555
|
+
currency TEXT NOT NULL CHECK (
|
|
6556
|
+
(meter = 'tokens.total' AND currency = '')
|
|
6557
|
+
OR (meter = 'cost.catalogEstimated' AND currency = 'USD')
|
|
6558
|
+
),
|
|
6559
|
+
limit_amount NUMERIC(78, 0) NOT NULL CHECK (limit_amount >= 0),
|
|
6560
|
+
period_kind TEXT NOT NULL CHECK (period_kind = 'calendarMonth'),
|
|
6561
|
+
enabled BOOLEAN NOT NULL,
|
|
6562
|
+
created_at TIMESTAMPTZ NOT NULL,
|
|
6563
|
+
updated_at TIMESTAMPTZ NOT NULL,
|
|
6564
|
+
PRIMARY KEY (project_id, id),
|
|
6565
|
+
UNIQUE (project_id, subject_type, subject_id, meter, currency)
|
|
6566
|
+
);
|
|
6567
|
+
|
|
6568
|
+
CREATE INDEX idx_ai_usage_limits_project_enabled
|
|
6569
|
+
ON ai_usage_limits (project_id, enabled, subject_type, subject_id, meter, currency);
|
|
6570
|
+
|
|
6571
|
+
CREATE TABLE ai_usage_limit_periods (
|
|
6572
|
+
project_id TEXT NOT NULL CHECK (length(btrim(project_id)) > 0),
|
|
6573
|
+
subject_type TEXT NOT NULL CHECK (
|
|
6574
|
+
subject_type IN ('project', 'group', 'user', 'serviceAccount')
|
|
6575
|
+
),
|
|
6576
|
+
subject_id TEXT NOT NULL CHECK (
|
|
6577
|
+
(subject_type = 'project' AND subject_id = '')
|
|
6578
|
+
OR (subject_type <> 'project' AND length(btrim(subject_id)) > 0)
|
|
6579
|
+
),
|
|
6580
|
+
meter TEXT NOT NULL CHECK (meter IN ('tokens.total', 'cost.catalogEstimated')),
|
|
6581
|
+
currency TEXT NOT NULL CHECK (
|
|
6582
|
+
(meter = 'tokens.total' AND currency = '')
|
|
6583
|
+
OR (meter = 'cost.catalogEstimated' AND currency = 'USD')
|
|
6584
|
+
),
|
|
6585
|
+
period_start TIMESTAMPTZ NOT NULL,
|
|
6586
|
+
period_end TIMESTAMPTZ NOT NULL,
|
|
6587
|
+
actual_amount NUMERIC(78, 0) NOT NULL DEFAULT 0 CHECK (actual_amount >= 0),
|
|
6588
|
+
reserved_amount NUMERIC(78, 0) NOT NULL DEFAULT 0 CHECK (reserved_amount >= 0),
|
|
6589
|
+
unknown_amount NUMERIC(78, 0) NOT NULL DEFAULT 0 CHECK (unknown_amount >= 0),
|
|
6590
|
+
accounting_status TEXT NOT NULL DEFAULT 'complete' CHECK (
|
|
6591
|
+
accounting_status IN ('complete', 'unavailable')
|
|
6592
|
+
),
|
|
6593
|
+
updated_at TIMESTAMPTZ NOT NULL,
|
|
6594
|
+
PRIMARY KEY (
|
|
6595
|
+
project_id, subject_type, subject_id, meter, currency, period_start
|
|
6596
|
+
)
|
|
6597
|
+
);
|
|
6598
|
+
|
|
6599
|
+
CREATE INDEX idx_ai_usage_limit_periods_subject_period
|
|
6600
|
+
ON ai_usage_limit_periods (
|
|
6601
|
+
project_id, subject_type, subject_id, period_start, period_end, meter, currency
|
|
6602
|
+
);
|
|
6603
|
+
|
|
6604
|
+
CREATE TABLE ai_model_call_reservations (
|
|
6605
|
+
project_id TEXT NOT NULL CHECK (length(btrim(project_id)) > 0),
|
|
6606
|
+
execution_id TEXT NOT NULL CHECK (length(btrim(execution_id)) > 0),
|
|
6607
|
+
attempt BIGINT NOT NULL CHECK (attempt >= 1),
|
|
6608
|
+
call_id TEXT NOT NULL CHECK (length(btrim(call_id)) > 0),
|
|
6609
|
+
buckets JSONB NOT NULL CHECK (jsonb_typeof(buckets) = 'array'),
|
|
6610
|
+
request_key TEXT NOT NULL,
|
|
6611
|
+
period_start TIMESTAMPTZ NOT NULL,
|
|
6612
|
+
period_end TIMESTAMPTZ NOT NULL,
|
|
6613
|
+
state TEXT NOT NULL CHECK (state IN ('active', 'reconciled', 'unknown')),
|
|
6614
|
+
usage_record_id TEXT,
|
|
6615
|
+
reserved_at TIMESTAMPTZ NOT NULL,
|
|
6616
|
+
updated_at TIMESTAMPTZ NOT NULL,
|
|
6617
|
+
PRIMARY KEY (project_id, execution_id, attempt, call_id),
|
|
6618
|
+
FOREIGN KEY (project_id, execution_id)
|
|
6619
|
+
REFERENCES executions(project_id, id)
|
|
6620
|
+
ON DELETE RESTRICT,
|
|
6621
|
+
FOREIGN KEY (project_id, usage_record_id)
|
|
6622
|
+
REFERENCES ai_model_call_usage(project_id, id)
|
|
6623
|
+
ON DELETE RESTRICT,
|
|
6624
|
+
CHECK (
|
|
6625
|
+
(
|
|
6626
|
+
state = 'reconciled'
|
|
6627
|
+
AND usage_record_id IS NOT NULL
|
|
6628
|
+
)
|
|
6629
|
+
OR (
|
|
6630
|
+
state <> 'reconciled'
|
|
6631
|
+
AND usage_record_id IS NULL
|
|
6632
|
+
)
|
|
6633
|
+
)
|
|
6634
|
+
);
|
|
6635
|
+
|
|
6636
|
+
CREATE INDEX idx_ai_model_call_reservations_active
|
|
6637
|
+
ON ai_model_call_reservations (project_id, state, period_end, execution_id, attempt, call_id);
|
|
6638
|
+
`;
|
|
6639
|
+
|
|
6640
|
+
// src/migrations/031-subagent-runs.sql
|
|
6641
|
+
var _031_subagent_runs_default = `-- Agent executions may either use a managed service account or inherit their parent's exact
|
|
6642
|
+
-- credentialed-user / auth-disabled authority.
|
|
6643
|
+
ALTER TABLE executions
|
|
6644
|
+
DROP CONSTRAINT executions_check3,
|
|
6645
|
+
ADD CONSTRAINT executions_agent_authority_check CHECK (
|
|
6646
|
+
(
|
|
6647
|
+
executor_kind = 'request'
|
|
6648
|
+
AND source_kind = 'http'
|
|
6649
|
+
AND executor_id = source_id
|
|
6650
|
+
AND (
|
|
6651
|
+
(
|
|
6652
|
+
authority_kind = 'principal'
|
|
6653
|
+
AND requested_by_user_id IS NOT DISTINCT FROM authority_user_id
|
|
6654
|
+
AND requested_by_service_account_id IS NOT DISTINCT FROM authority_service_account_id
|
|
6655
|
+
)
|
|
6656
|
+
OR (
|
|
6657
|
+
authority_kind = 'disabled'
|
|
6658
|
+
AND requested_by_user_id IS NULL
|
|
6659
|
+
AND requested_by_service_account_id IS NULL
|
|
6660
|
+
)
|
|
6661
|
+
)
|
|
6662
|
+
)
|
|
6663
|
+
OR (
|
|
6664
|
+
executor_kind IN ('action', 'pipeline', 'projection', 'rule', 'sync', 'webhook', 'workflow')
|
|
6665
|
+
AND authority_kind = 'trustedPrimitive'
|
|
6666
|
+
AND executor_kind = authority_primitive_kind
|
|
6667
|
+
)
|
|
6668
|
+
OR (
|
|
6669
|
+
executor_kind = 'agent'
|
|
6670
|
+
AND source_kind = 'execution'
|
|
6671
|
+
AND (
|
|
6672
|
+
(
|
|
6673
|
+
authority_kind = 'principal'
|
|
6674
|
+
AND authority_service_account_id IS NOT NULL
|
|
6675
|
+
AND authority_session_id IS NULL
|
|
6676
|
+
AND authority_access_token_id IS NULL
|
|
6677
|
+
)
|
|
6678
|
+
OR (
|
|
6679
|
+
authority_kind = 'principal'
|
|
6680
|
+
AND authority_user_id IS NOT NULL
|
|
6681
|
+
AND num_nonnulls(authority_session_id, authority_access_token_id) = 1
|
|
6682
|
+
)
|
|
6683
|
+
OR authority_kind = 'disabled'
|
|
6684
|
+
)
|
|
6685
|
+
)
|
|
6686
|
+
OR (executor_kind = 'kernel' AND authority_kind = 'kernel')
|
|
6687
|
+
);
|
|
6688
|
+
|
|
6689
|
+
ALTER TABLE agent_runs
|
|
6690
|
+
ADD COLUMN kind TEXT NOT NULL DEFAULT 'conversation'
|
|
6691
|
+
CHECK (kind IN ('conversation', 'subagent')),
|
|
6692
|
+
ADD COLUMN parent_run_id TEXT,
|
|
6693
|
+
ADD COLUMN spawn_key TEXT,
|
|
6694
|
+
ADD COLUMN spec JSONB,
|
|
6695
|
+
ADD COLUMN result JSONB,
|
|
6696
|
+
ALTER COLUMN thread_id DROP NOT NULL,
|
|
6697
|
+
ALTER COLUMN agent_id DROP NOT NULL,
|
|
6698
|
+
ALTER COLUMN trigger_message_id DROP NOT NULL,
|
|
6699
|
+
ADD CONSTRAINT fk_agent_runs_parent
|
|
6700
|
+
FOREIGN KEY (project_id, parent_run_id)
|
|
6701
|
+
REFERENCES agent_runs (project_id, id)
|
|
6702
|
+
ON DELETE RESTRICT,
|
|
6703
|
+
ADD CONSTRAINT agent_runs_kind_fields CHECK (
|
|
6704
|
+
(
|
|
6705
|
+
kind = 'conversation'
|
|
6706
|
+
AND thread_id IS NOT NULL
|
|
6707
|
+
AND agent_id IS NOT NULL
|
|
6708
|
+
AND trigger_message_id IS NOT NULL
|
|
6709
|
+
AND parent_run_id IS NULL
|
|
6710
|
+
AND spawn_key IS NULL
|
|
6711
|
+
AND spec IS NULL
|
|
6712
|
+
AND result IS NULL
|
|
6713
|
+
)
|
|
6714
|
+
OR
|
|
6715
|
+
(
|
|
6716
|
+
kind = 'subagent'
|
|
6717
|
+
AND thread_id IS NULL
|
|
6718
|
+
AND agent_id IS NULL
|
|
6719
|
+
AND trigger_message_id IS NULL
|
|
6720
|
+
AND parent_run_id IS NOT NULL
|
|
6721
|
+
AND spawn_key IS NOT NULL
|
|
6722
|
+
AND spec IS NOT NULL
|
|
6723
|
+
)
|
|
6724
|
+
),
|
|
6725
|
+
ADD CONSTRAINT agent_runs_subagent_result CHECK (
|
|
6726
|
+
kind <> 'subagent' OR status <> 'succeeded' OR result IS NOT NULL
|
|
6727
|
+
),
|
|
6728
|
+
ADD CONSTRAINT uq_agent_runs_parent_spawn_key UNIQUE (project_id, parent_run_id, spawn_key);
|
|
6729
|
+
|
|
6730
|
+
CREATE INDEX idx_agent_runs_project_parent_started
|
|
6731
|
+
ON agent_runs (
|
|
6732
|
+
project_id,
|
|
6733
|
+
parent_run_id,
|
|
6734
|
+
status,
|
|
6735
|
+
COALESCE(started_at, created_at) DESC,
|
|
6736
|
+
id DESC
|
|
6737
|
+
);
|
|
6738
|
+
`;
|
|
6739
|
+
|
|
6740
|
+
// src/migrations/032-conversation-run-spec.sql
|
|
6741
|
+
var _032_conversation_run_spec_default = `ALTER TABLE agent_runs
|
|
6742
|
+
DROP CONSTRAINT agent_runs_kind_fields,
|
|
6743
|
+
ADD CONSTRAINT agent_runs_kind_fields CHECK (
|
|
6744
|
+
(
|
|
6745
|
+
kind = 'conversation'
|
|
6746
|
+
AND thread_id IS NOT NULL
|
|
6747
|
+
AND agent_id IS NOT NULL
|
|
6748
|
+
AND trigger_message_id IS NOT NULL
|
|
6749
|
+
AND parent_run_id IS NULL
|
|
6750
|
+
AND spawn_key IS NULL
|
|
6751
|
+
AND result IS NULL
|
|
6752
|
+
)
|
|
6753
|
+
OR
|
|
6754
|
+
(
|
|
6755
|
+
kind = 'subagent'
|
|
6756
|
+
AND thread_id IS NULL
|
|
6757
|
+
AND agent_id IS NULL
|
|
6758
|
+
AND trigger_message_id IS NULL
|
|
6759
|
+
AND parent_run_id IS NOT NULL
|
|
6760
|
+
AND spawn_key IS NOT NULL
|
|
6761
|
+
AND spec IS NOT NULL
|
|
6762
|
+
)
|
|
6763
|
+
);
|
|
6764
|
+
`;
|
|
6765
|
+
|
|
6766
|
+
// src/migrations/033-retire-agent-definitions.sql
|
|
6767
|
+
var _033_retire_agent_definitions_default = `DROP INDEX IF EXISTS idx_agent_threads_project_agent_activity;
|
|
6768
|
+
ALTER TABLE workflow_agent_node_runs RENAME COLUMN agent_id TO actor_id;
|
|
6769
|
+
DROP INDEX IF EXISTS idx_agent_runs_project_agent_started;
|
|
6770
|
+
|
|
6771
|
+
ALTER TABLE agent_threads DROP COLUMN agent_id;
|
|
6772
|
+
|
|
6773
|
+
ALTER TABLE agent_runs
|
|
6774
|
+
DROP CONSTRAINT agent_runs_kind_fields,
|
|
6775
|
+
DROP COLUMN agent_id,
|
|
6776
|
+
ADD CONSTRAINT agent_runs_kind_fields CHECK (
|
|
6777
|
+
(
|
|
6778
|
+
kind = 'conversation'
|
|
6779
|
+
AND thread_id IS NOT NULL
|
|
6780
|
+
AND trigger_message_id IS NOT NULL
|
|
6781
|
+
AND parent_run_id IS NULL
|
|
6782
|
+
AND spawn_key IS NULL
|
|
6783
|
+
AND result IS NULL
|
|
6784
|
+
)
|
|
6785
|
+
OR
|
|
6786
|
+
(
|
|
6787
|
+
kind = 'subagent'
|
|
6788
|
+
AND thread_id IS NULL
|
|
6789
|
+
AND trigger_message_id IS NULL
|
|
6790
|
+
AND parent_run_id IS NOT NULL
|
|
6791
|
+
AND spawn_key IS NOT NULL
|
|
6792
|
+
AND spec IS NOT NULL
|
|
6793
|
+
)
|
|
6794
|
+
);
|
|
6795
|
+
`;
|
|
6796
|
+
|
|
6797
|
+
// src/migrations/034-device-authorizations.sql
|
|
6798
|
+
var _034_device_authorizations_default = `CREATE TABLE auth_device_authorizations (
|
|
6799
|
+
project_id TEXT NOT NULL,
|
|
6800
|
+
id TEXT NOT NULL,
|
|
6801
|
+
device_code_hash TEXT NOT NULL,
|
|
6802
|
+
user_code TEXT NOT NULL,
|
|
6803
|
+
client_name TEXT NOT NULL,
|
|
6804
|
+
token_name TEXT NOT NULL,
|
|
6805
|
+
token_expires_at TIMESTAMPTZ NOT NULL,
|
|
6806
|
+
status TEXT NOT NULL CHECK (status IN ('pending', 'approved', 'denied', 'consumed')),
|
|
6807
|
+
approved_user_id TEXT,
|
|
6808
|
+
approved_session_id TEXT,
|
|
6809
|
+
created_at TIMESTAMPTZ NOT NULL,
|
|
6810
|
+
expires_at TIMESTAMPTZ NOT NULL,
|
|
6811
|
+
approved_at TIMESTAMPTZ,
|
|
6812
|
+
denied_at TIMESTAMPTZ,
|
|
6813
|
+
consumed_at TIMESTAMPTZ,
|
|
6814
|
+
PRIMARY KEY (project_id, id),
|
|
6815
|
+
UNIQUE (project_id, user_code),
|
|
6816
|
+
FOREIGN KEY (project_id, approved_user_id) REFERENCES auth_users (project_id, id),
|
|
6817
|
+
FOREIGN KEY (project_id, approved_session_id) REFERENCES auth_sessions (project_id, id)
|
|
6818
|
+
);
|
|
6819
|
+
|
|
6820
|
+
CREATE INDEX idx_auth_device_authorizations_pending
|
|
6821
|
+
ON auth_device_authorizations (project_id, status, expires_at DESC);
|
|
6822
|
+
`;
|
|
6823
|
+
|
|
6824
|
+
// src/migrations.ts
|
|
6825
|
+
var POSTGRES_STORAGE_ADAPTER_ID = "SixbPostgresStorage";
|
|
6826
|
+
function pgSql(id, sqlText) {
|
|
6827
|
+
return pgStep(id, (context) => context.exec(sqlText), { checksum: checksum(sqlText) });
|
|
6828
|
+
}
|
|
6829
|
+
function pgStep(id, up, options = {}) {
|
|
6830
|
+
return step(id, up, options);
|
|
6831
|
+
}
|
|
6832
|
+
function createPostgresStorageMigrators(sql, schemaName) {
|
|
6833
|
+
return [
|
|
6834
|
+
createPostgresMigrator({
|
|
6835
|
+
sql,
|
|
6836
|
+
schemaName,
|
|
6837
|
+
migrations: postgresStorageMigrations
|
|
6838
|
+
})
|
|
6839
|
+
];
|
|
6840
|
+
}
|
|
6841
|
+
function createPostgresMigrator(params) {
|
|
6842
|
+
return {
|
|
6843
|
+
adapterId: params.migrations.adapterId,
|
|
6844
|
+
latestVersion: params.migrations.latestVersion,
|
|
6845
|
+
async status() {
|
|
6846
|
+
return describeMigrationHistory({
|
|
6847
|
+
migrations: params.migrations,
|
|
6848
|
+
rows: await readPostgresHistory(params.sql, params.schemaName, params.migrations.adapterId)
|
|
6849
|
+
});
|
|
6850
|
+
},
|
|
6851
|
+
migrate() {
|
|
6852
|
+
return withPostgresMigrationLock(params, async (sql) => {
|
|
6853
|
+
const session = postgresMigrationSession(sql, params.schemaName);
|
|
6854
|
+
return runMigrationSet({
|
|
6855
|
+
context: session.context,
|
|
6856
|
+
migrations: params.migrations,
|
|
6857
|
+
state: session.state
|
|
6858
|
+
});
|
|
6859
|
+
});
|
|
6860
|
+
}
|
|
6861
|
+
};
|
|
6862
|
+
}
|
|
6863
|
+
async function migratePostgresStorage(sql, schemaName) {
|
|
6864
|
+
await createPostgresMigrator({
|
|
6865
|
+
sql,
|
|
6866
|
+
schemaName,
|
|
6867
|
+
migrations: postgresStorageMigrations
|
|
6868
|
+
}).migrate();
|
|
6869
|
+
}
|
|
6870
|
+
async function dropSchema(sql, schemaName) {
|
|
6871
|
+
await sql.unsafe(`DROP SCHEMA IF EXISTS ${quoteIdent(schemaName)} CASCADE`);
|
|
6872
|
+
}
|
|
6873
|
+
function quoteIdent(identifier) {
|
|
6874
|
+
if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(identifier)) {
|
|
6875
|
+
throw new Error(`[SixbPg] Invalid identifier: ${identifier}`);
|
|
6876
|
+
}
|
|
6877
|
+
return `"${identifier}"`;
|
|
6878
|
+
}
|
|
6879
|
+
function postgresMigrationSession(sql, schemaName) {
|
|
6880
|
+
const schema = quoteIdent(schemaName);
|
|
6881
|
+
const migrationsTable = `${schema}.sixb_migrations`;
|
|
6882
|
+
let active = null;
|
|
6883
|
+
const connection = () => active ?? sql;
|
|
6884
|
+
return {
|
|
6885
|
+
context: {
|
|
6886
|
+
async exec(sqlText) {
|
|
6887
|
+
await connection().unsafe(sqlText);
|
|
6888
|
+
}
|
|
6889
|
+
},
|
|
6890
|
+
state: {
|
|
6891
|
+
async ensure() {
|
|
6892
|
+
await sql.unsafe(`
|
|
6893
|
+
CREATE SCHEMA IF NOT EXISTS ${schema};
|
|
6894
|
+
|
|
6895
|
+
CREATE TABLE IF NOT EXISTS ${migrationsTable} (
|
|
6896
|
+
adapter_id TEXT NOT NULL,
|
|
6897
|
+
version INTEGER NOT NULL,
|
|
6898
|
+
id TEXT NOT NULL,
|
|
6899
|
+
checksum TEXT,
|
|
6900
|
+
status TEXT NOT NULL CHECK (status IN ('started', 'applied')),
|
|
6901
|
+
started_at TIMESTAMPTZ NOT NULL,
|
|
6902
|
+
finished_at TIMESTAMPTZ,
|
|
6903
|
+
PRIMARY KEY (adapter_id, version)
|
|
6904
|
+
);
|
|
6905
|
+
`);
|
|
6906
|
+
},
|
|
6907
|
+
async readHistory(adapterId) {
|
|
6908
|
+
const rows = await sql.unsafe(`SELECT * FROM ${migrationsTable} WHERE adapter_id = $1 ORDER BY version`, [adapterId]);
|
|
6909
|
+
return rows.map(rowToMigrationRecord);
|
|
6910
|
+
},
|
|
6911
|
+
async markStarted(adapterId, migration, at) {
|
|
6912
|
+
await connection().unsafe(`
|
|
6913
|
+
INSERT INTO ${migrationsTable} (
|
|
6914
|
+
adapter_id, version, id, checksum, status, started_at, finished_at
|
|
6915
|
+
) VALUES ($1, $2, $3, $4, 'started', $5, NULL)
|
|
6916
|
+
ON CONFLICT(adapter_id, version) DO UPDATE SET
|
|
6917
|
+
id = excluded.id,
|
|
6918
|
+
checksum = excluded.checksum,
|
|
6919
|
+
status = excluded.status,
|
|
6920
|
+
started_at = excluded.started_at,
|
|
6921
|
+
finished_at = NULL
|
|
6922
|
+
`, [adapterId, migration.version, migration.id, migration.checksum ?? null, at]);
|
|
6923
|
+
},
|
|
6924
|
+
async markApplied(adapterId, migration, at) {
|
|
6925
|
+
await connection().unsafe(`
|
|
6926
|
+
UPDATE ${migrationsTable}
|
|
6927
|
+
SET status = 'applied', finished_at = $1
|
|
6928
|
+
WHERE adapter_id = $2 AND version = $3
|
|
6929
|
+
`, [at, adapterId, migration.version]);
|
|
6930
|
+
},
|
|
6931
|
+
async transaction(run) {
|
|
6932
|
+
const previous = active;
|
|
6933
|
+
active = sql;
|
|
6934
|
+
await sql.unsafe("BEGIN");
|
|
6935
|
+
try {
|
|
6936
|
+
await sql.unsafe(`SET LOCAL search_path TO ${schema}`);
|
|
6937
|
+
const result = await run();
|
|
6938
|
+
await sql.unsafe("COMMIT");
|
|
6363
6939
|
return result;
|
|
6364
6940
|
} catch (error) {
|
|
6365
6941
|
await sql.unsafe("ROLLBACK");
|
|
@@ -6447,7 +7023,13 @@ var postgresStorageMigrations = defineMigrations({
|
|
|
6447
7023
|
pgSql("025-connector-connections", _025_connector_connections_default),
|
|
6448
7024
|
pgSql("026-ai-cost-accounting", _026_ai_cost_accounting_default),
|
|
6449
7025
|
pgSql("027-agent-context-checkpoints", _027_agent_context_checkpoints_default),
|
|
6450
|
-
pgSql("028-object-override-edit-times", _028_object_override_edit_times_default)
|
|
7026
|
+
pgSql("028-object-override-edit-times", _028_object_override_edit_times_default),
|
|
7027
|
+
pgSql("029-model-accounting", _029_model_accounting_default),
|
|
7028
|
+
pgSql("030-ai-usage-limits", _030_ai_usage_limits_default),
|
|
7029
|
+
pgSql("031-subagent-runs", _031_subagent_runs_default),
|
|
7030
|
+
pgSql("032-conversation-run-spec", _032_conversation_run_spec_default),
|
|
7031
|
+
pgSql("033-retire-agent-definitions", _033_retire_agent_definitions_default),
|
|
7032
|
+
pgSql("034-device-authorizations", _034_device_authorizations_default)
|
|
6451
7033
|
]
|
|
6452
7034
|
});
|
|
6453
7035
|
|
|
@@ -10165,32 +10747,366 @@ import {
|
|
|
10165
10747
|
aiModelCallCostMatchesUsage,
|
|
10166
10748
|
buildAiAccountingOverviewFromFragments,
|
|
10167
10749
|
normalizeAiAccountingQuery,
|
|
10168
|
-
normalizeAiModelCallAccountingQuery,
|
|
10750
|
+
normalizeAiModelCallAccountingQuery as normalizeAiModelCallAccountingQuery2,
|
|
10169
10751
|
normalizeAiModelCallCostRecord,
|
|
10170
10752
|
parseAiModelCallCostDetails,
|
|
10171
10753
|
toAccountingItem
|
|
10172
10754
|
} from "@sixb/core/internal/ai-cost-storage-provider";
|
|
10755
|
+
import { normalizeModelProviderIds as normalizeModelProviderIds2 } from "@sixb/core/models";
|
|
10173
10756
|
import {
|
|
10174
10757
|
AiCostStorageError
|
|
10175
10758
|
} from "@sixb/core/storage";
|
|
10176
|
-
|
|
10177
|
-
|
|
10178
|
-
|
|
10179
|
-
|
|
10180
|
-
|
|
10181
|
-
|
|
10182
|
-
|
|
10183
|
-
|
|
10184
|
-
|
|
10185
|
-
|
|
10186
|
-
|
|
10187
|
-
|
|
10188
|
-
|
|
10189
|
-
|
|
10190
|
-
|
|
10191
|
-
|
|
10192
|
-
|
|
10193
|
-
|
|
10759
|
+
|
|
10760
|
+
// src/ai-cost-groups.ts
|
|
10761
|
+
import {
|
|
10762
|
+
aiModelCallGroupFragmentFromRow,
|
|
10763
|
+
buildAiModelCallGroupsFromFragments,
|
|
10764
|
+
normalizeAiModelCallAccountingQuery
|
|
10765
|
+
} from "@sixb/core/internal/ai-cost-storage-provider";
|
|
10766
|
+
async function listAiModelCallGroups(sql, input) {
|
|
10767
|
+
const query = normalizeAiModelCallAccountingQuery(input);
|
|
10768
|
+
const parameters = [
|
|
10769
|
+
query.projectId,
|
|
10770
|
+
query.from.toISOString(),
|
|
10771
|
+
query.to.toISOString(),
|
|
10772
|
+
query.providerId ?? null,
|
|
10773
|
+
query.modelId ?? null,
|
|
10774
|
+
query.valuationStatus ?? null
|
|
10775
|
+
];
|
|
10776
|
+
const [count] = await sql.unsafe(COUNT_SQL, parameters);
|
|
10777
|
+
const rows = await sql.unsafe(PAGE_SQL, [
|
|
10778
|
+
...parameters,
|
|
10779
|
+
query.limit,
|
|
10780
|
+
query.offset
|
|
10781
|
+
]);
|
|
10782
|
+
const total = Number(count?.total);
|
|
10783
|
+
if (!Number.isSafeInteger(total) || total < 0)
|
|
10784
|
+
throw new RangeError("[Sixb] Invalid model-call group count.");
|
|
10785
|
+
const items = buildAiModelCallGroupsFromFragments(rows.map(aiModelCallGroupFragmentFromRow));
|
|
10786
|
+
return { items, total, hasMore: query.offset + items.length < total };
|
|
10787
|
+
}
|
|
10788
|
+
var FILTERED_SQL = `
|
|
10789
|
+
WITH filtered AS (
|
|
10790
|
+
SELECT COALESCE(root_agent.execution_id, usage.execution_id) AS root_execution_id,
|
|
10791
|
+
usage.execution_id, usage.provider_id, usage.requested_model_id, usage.occurred_at,
|
|
10792
|
+
usage.total_tokens,
|
|
10793
|
+
COALESCE(cost.status, 'unvalued') AS valuation_status,
|
|
10794
|
+
cost.currency AS amount_currency, cost.amount_nanos,
|
|
10795
|
+
root_agent.id AS root_agent_run_id,
|
|
10796
|
+
root_agent.thread_id AS root_thread_id,
|
|
10797
|
+
direct_agent.kind AS direct_kind,
|
|
10798
|
+
direct_agent.id AS direct_run_id,
|
|
10799
|
+
direct_agent.parent_run_id AS parent_run_id,
|
|
10800
|
+
workflow_node.node_id AS workflow_agent_step_id,
|
|
10801
|
+
workflow_agent.node_run_id AS workflow_node_run_id,
|
|
10802
|
+
workflow_node.workflow_id AS workflow_id,
|
|
10803
|
+
workflow_node.workflow_run_id AS workflow_run_id
|
|
10804
|
+
FROM ai_model_call_usage AS usage
|
|
10805
|
+
LEFT JOIN agent_runs AS direct_agent
|
|
10806
|
+
ON direct_agent.project_id = usage.project_id AND direct_agent.execution_id = usage.execution_id
|
|
10807
|
+
LEFT JOIN agent_runs AS root_agent
|
|
10808
|
+
ON root_agent.project_id = usage.project_id
|
|
10809
|
+
AND root_agent.id = COALESCE(direct_agent.parent_run_id, direct_agent.id)
|
|
10810
|
+
LEFT JOIN workflow_agent_node_runs AS workflow_agent
|
|
10811
|
+
ON workflow_agent.project_id = usage.project_id AND workflow_agent.execution_id = usage.execution_id
|
|
10812
|
+
LEFT JOIN workflow_node_runs AS workflow_node
|
|
10813
|
+
ON workflow_node.project_id = workflow_agent.project_id AND workflow_node.id = workflow_agent.node_run_id
|
|
10814
|
+
LEFT JOIN ai_model_call_valuations AS cost
|
|
10815
|
+
ON cost.project_id = usage.project_id AND cost.usage_record_id = usage.id
|
|
10816
|
+
WHERE usage.project_id = $1
|
|
10817
|
+
AND usage.occurred_at >= $2::timestamptz
|
|
10818
|
+
AND usage.occurred_at < $3::timestamptz
|
|
10819
|
+
AND ($4::text IS NULL OR usage.provider_id = $4)
|
|
10820
|
+
AND ($5::text IS NULL OR usage.requested_model_id = $5)
|
|
10821
|
+
AND ($6::text IS NULL OR COALESCE(cost.status, 'unvalued') = $6)
|
|
10822
|
+
), roots AS (
|
|
10823
|
+
SELECT root_execution_id, MAX(occurred_at) AS last_call_at
|
|
10824
|
+
FROM filtered GROUP BY root_execution_id
|
|
10825
|
+
)
|
|
10826
|
+
`;
|
|
10827
|
+
var COUNT_SQL = `${FILTERED_SQL} SELECT CAST(COUNT(*) AS TEXT) AS total FROM roots`;
|
|
10828
|
+
var PAGE_SQL = `${FILTERED_SQL}, page AS (
|
|
10829
|
+
SELECT root_execution_id FROM roots
|
|
10830
|
+
ORDER BY last_call_at DESC, root_execution_id COLLATE "C" ASC LIMIT $7 OFFSET $8
|
|
10831
|
+
)
|
|
10832
|
+
SELECT filtered.root_execution_id, execution_id, provider_id, requested_model_id,
|
|
10833
|
+
MIN(occurred_at) AS first_call_at, MAX(occurred_at) AS last_call_at,
|
|
10834
|
+
CAST(COUNT(*) AS TEXT) AS call_count,
|
|
10835
|
+
CAST(COUNT(total_tokens) AS TEXT) AS token_count,
|
|
10836
|
+
CAST(SUM(total_tokens) AS TEXT) AS total_tokens,
|
|
10837
|
+
valuation_status, amount_currency,
|
|
10838
|
+
CAST(SUM(amount_nanos / 1000000000000) AS TEXT) AS amount_high,
|
|
10839
|
+
CAST(SUM((amount_nanos / 1000000) % 1000000) AS TEXT) AS amount_middle,
|
|
10840
|
+
CAST(SUM(amount_nanos % 1000000) AS TEXT) AS amount_low,
|
|
10841
|
+
MAX(root_agent_run_id) AS root_agent_run_id,
|
|
10842
|
+
MAX(root_thread_id) AS root_thread_id,
|
|
10843
|
+
MAX(direct_kind) AS direct_kind,
|
|
10844
|
+
MAX(direct_run_id) AS direct_run_id,
|
|
10845
|
+
MAX(parent_run_id) AS parent_run_id,
|
|
10846
|
+
MAX(workflow_agent_step_id) AS workflow_agent_step_id,
|
|
10847
|
+
MAX(workflow_node_run_id) AS workflow_node_run_id,
|
|
10848
|
+
MAX(workflow_id) AS workflow_id,
|
|
10849
|
+
MAX(workflow_run_id) AS workflow_run_id
|
|
10850
|
+
FROM filtered JOIN page ON page.root_execution_id = filtered.root_execution_id
|
|
10851
|
+
GROUP BY filtered.root_execution_id, execution_id, provider_id, requested_model_id,
|
|
10852
|
+
valuation_status, amount_currency
|
|
10853
|
+
`;
|
|
10854
|
+
|
|
10855
|
+
// src/pg-ai-usage-storage.ts
|
|
10856
|
+
import { isModelReasoning, normalizeModelProviderIds } from "@sixb/core/models";
|
|
10857
|
+
import {
|
|
10858
|
+
AiUsageStorageError,
|
|
10859
|
+
aggregateAiModelCallUsage,
|
|
10860
|
+
assertAiUsageExecutionId,
|
|
10861
|
+
normalizeAiModelCallRecord
|
|
10862
|
+
} from "@sixb/core/storage";
|
|
10863
|
+
class PgAiUsageStorage {
|
|
10864
|
+
sql;
|
|
10865
|
+
constructor(sql) {
|
|
10866
|
+
this.sql = sql;
|
|
10867
|
+
}
|
|
10868
|
+
async recordModelCall(input) {
|
|
10869
|
+
const record = normalizeAiModelCallRecord(input);
|
|
10870
|
+
return runPgTransaction(this.sql, async (tx) => {
|
|
10871
|
+
const existing = await this.findByIdempotencyKey(tx, record);
|
|
10872
|
+
if (existing)
|
|
10873
|
+
return { record: existing, created: false };
|
|
10874
|
+
await this.assertExecutionExists(tx, record.projectId, record.executionId);
|
|
10875
|
+
const inserted = await tx`
|
|
10876
|
+
INSERT INTO ai_model_call_usage (
|
|
10877
|
+
project_id,
|
|
10878
|
+
id,
|
|
10879
|
+
execution_id,
|
|
10880
|
+
attempt,
|
|
10881
|
+
call_id,
|
|
10882
|
+
provider_id,
|
|
10883
|
+
provider_ids,
|
|
10884
|
+
requested_model_id,
|
|
10885
|
+
requested_reasoning,
|
|
10886
|
+
response_model_id,
|
|
10887
|
+
response_id,
|
|
10888
|
+
input_tokens,
|
|
10889
|
+
output_tokens,
|
|
10890
|
+
total_tokens,
|
|
10891
|
+
uncached_input_tokens,
|
|
10892
|
+
cache_read_input_tokens,
|
|
10893
|
+
cache_write_input_tokens,
|
|
10894
|
+
text_output_tokens,
|
|
10895
|
+
reasoning_output_tokens,
|
|
10896
|
+
reporting_status,
|
|
10897
|
+
raw_usage,
|
|
10898
|
+
occurred_at,
|
|
10899
|
+
recorded_at
|
|
10900
|
+
) VALUES (
|
|
10901
|
+
${record.projectId},
|
|
10902
|
+
${record.id},
|
|
10903
|
+
${record.executionId},
|
|
10904
|
+
${record.attempt},
|
|
10905
|
+
${record.callId},
|
|
10906
|
+
${record.providerId},
|
|
10907
|
+
${record.providerIds === undefined ? null : JSON.stringify(record.providerIds)}::jsonb,
|
|
10908
|
+
${record.requestedModelId},
|
|
10909
|
+
${record.requestedReasoning === undefined ? null : JSON.stringify(record.requestedReasoning)}::text::jsonb,
|
|
10910
|
+
${record.responseModelId ?? null},
|
|
10911
|
+
${record.responseId},
|
|
10912
|
+
${record.usage.inputTokens ?? null},
|
|
10913
|
+
${record.usage.outputTokens ?? null},
|
|
10914
|
+
${record.usage.totalTokens ?? null},
|
|
10915
|
+
${record.usage.uncachedInputTokens ?? null},
|
|
10916
|
+
${record.usage.cacheReadInputTokens ?? null},
|
|
10917
|
+
${record.usage.cacheWriteInputTokens ?? null},
|
|
10918
|
+
${record.usage.textOutputTokens ?? null},
|
|
10919
|
+
${record.usage.reasoningOutputTokens ?? null},
|
|
10920
|
+
${record.usage.reportingStatus},
|
|
10921
|
+
${record.rawUsage === undefined ? null : JSON.stringify(record.rawUsage)}::text::jsonb,
|
|
10922
|
+
${record.occurredAt},
|
|
10923
|
+
${record.recordedAt}
|
|
10924
|
+
)
|
|
10925
|
+
ON CONFLICT DO NOTHING
|
|
10926
|
+
RETURNING id
|
|
10927
|
+
`;
|
|
10928
|
+
if (inserted.length === 0) {
|
|
10929
|
+
const concurrent = await this.findByIdempotencyKey(tx, record);
|
|
10930
|
+
if (concurrent)
|
|
10931
|
+
return { record: concurrent, created: false };
|
|
10932
|
+
throw new AiUsageStorageError("duplicate_id", `[SixbPg] AI usage record '${record.id}' already exists for project '${record.projectId}'.`);
|
|
10933
|
+
}
|
|
10934
|
+
for (const groupId of record.requesterGroupIds) {
|
|
10935
|
+
await tx`
|
|
10936
|
+
INSERT INTO ai_model_call_usage_groups (
|
|
10937
|
+
project_id, usage_record_id, group_id, occurred_at
|
|
10938
|
+
) VALUES (
|
|
10939
|
+
${record.projectId}, ${record.id}, ${groupId}, ${record.occurredAt}
|
|
10940
|
+
)
|
|
10941
|
+
`;
|
|
10942
|
+
}
|
|
10943
|
+
return { record: await this.requireById(tx, record.projectId, record.id), created: true };
|
|
10944
|
+
});
|
|
10945
|
+
}
|
|
10946
|
+
async summarizeExecution(input) {
|
|
10947
|
+
const [summary] = await this.summarizeExecutions({
|
|
10948
|
+
projectId: input.projectId,
|
|
10949
|
+
executionIds: [input.executionId]
|
|
10950
|
+
});
|
|
10951
|
+
return summary;
|
|
10952
|
+
}
|
|
10953
|
+
async getLatestForExecution(input) {
|
|
10954
|
+
assertNonBlankProjectId(input.projectId);
|
|
10955
|
+
assertAiUsageExecutionId(input.executionId);
|
|
10956
|
+
const [row] = await this.sql`
|
|
10957
|
+
SELECT * FROM ai_model_call_usage
|
|
10958
|
+
WHERE project_id = ${input.projectId} AND execution_id = ${input.executionId}
|
|
10959
|
+
ORDER BY occurred_at DESC, id DESC
|
|
10960
|
+
LIMIT 1
|
|
10961
|
+
`;
|
|
10962
|
+
return row ? this.rowToRecord(this.sql, row) : null;
|
|
10963
|
+
}
|
|
10964
|
+
async summarizeExecutions(input) {
|
|
10965
|
+
assertNonBlankProjectId(input.projectId);
|
|
10966
|
+
for (const executionId of input.executionIds)
|
|
10967
|
+
assertAiUsageExecutionId(executionId);
|
|
10968
|
+
if (input.executionIds.length === 0)
|
|
10969
|
+
return [];
|
|
10970
|
+
const rows = await this.executionRows(this.sql, input);
|
|
10971
|
+
return aggregateExecutionRows(input.executionIds, rows);
|
|
10972
|
+
}
|
|
10973
|
+
async findByIdempotencyKey(sql, record) {
|
|
10974
|
+
const rows = await sql`
|
|
10975
|
+
SELECT * FROM ai_model_call_usage
|
|
10976
|
+
WHERE project_id = ${record.projectId}
|
|
10977
|
+
AND execution_id = ${record.executionId}
|
|
10978
|
+
AND attempt = ${record.attempt}
|
|
10979
|
+
AND call_id = ${record.callId}
|
|
10980
|
+
AND response_id = ${record.responseId}
|
|
10981
|
+
`;
|
|
10982
|
+
return rows[0] ? this.rowToRecord(sql, rows[0]) : null;
|
|
10983
|
+
}
|
|
10984
|
+
async requireById(sql, projectId, id) {
|
|
10985
|
+
const [row] = await sql`
|
|
10986
|
+
SELECT * FROM ai_model_call_usage WHERE project_id = ${projectId} AND id = ${id}
|
|
10987
|
+
`;
|
|
10988
|
+
if (!row)
|
|
10989
|
+
throw new Error(`[SixbPg] AI usage record '${id}' disappeared after insert.`);
|
|
10990
|
+
return this.rowToRecord(sql, row);
|
|
10991
|
+
}
|
|
10992
|
+
async rowToRecord(sql, row) {
|
|
10993
|
+
const groupRows = await sql`
|
|
10994
|
+
SELECT group_id FROM ai_model_call_usage_groups
|
|
10995
|
+
WHERE project_id = ${row.project_id} AND usage_record_id = ${row.id}
|
|
10996
|
+
ORDER BY group_id
|
|
10997
|
+
`;
|
|
10998
|
+
return {
|
|
10999
|
+
id: row.id,
|
|
11000
|
+
projectId: row.project_id,
|
|
11001
|
+
executionId: row.execution_id,
|
|
11002
|
+
attempt: Number(row.attempt),
|
|
11003
|
+
callId: row.call_id,
|
|
11004
|
+
requesterGroupIds: groupRows.map((group) => group.group_id),
|
|
11005
|
+
providerId: row.provider_id,
|
|
11006
|
+
...row.provider_ids === null ? {} : {
|
|
11007
|
+
providerIds: normalizeModelProviderIds(typeof row.provider_ids === "string" ? JSON.parse(row.provider_ids) : row.provider_ids)
|
|
11008
|
+
},
|
|
11009
|
+
requestedModelId: row.requested_model_id,
|
|
11010
|
+
...row.requested_reasoning === null ? {} : { requestedReasoning: requestedReasoningFromRow(row.requested_reasoning) },
|
|
11011
|
+
...row.response_model_id === null ? {} : { responseModelId: row.response_model_id },
|
|
11012
|
+
responseId: row.response_id,
|
|
11013
|
+
usage: {
|
|
11014
|
+
...usageFromRow(row),
|
|
11015
|
+
...row.total_tokens === null ? {} : { totalTokens: Number(row.total_tokens) },
|
|
11016
|
+
reportingStatus: row.reporting_status
|
|
11017
|
+
},
|
|
11018
|
+
...row.raw_usage === null ? {} : { rawUsage: rawUsageFromRow(row.raw_usage) },
|
|
11019
|
+
occurredAt: new Date(row.occurred_at),
|
|
11020
|
+
recordedAt: new Date(row.recorded_at)
|
|
11021
|
+
};
|
|
11022
|
+
}
|
|
11023
|
+
executionRows(sql, input) {
|
|
11024
|
+
const executionIds = [...new Set(input.executionIds)];
|
|
11025
|
+
return sql`
|
|
11026
|
+
WITH requested(execution_id) AS (
|
|
11027
|
+
SELECT * FROM unnest(${sql.array(executionIds)}::text[])
|
|
11028
|
+
)
|
|
11029
|
+
SELECT usage.* FROM requested
|
|
11030
|
+
JOIN ai_model_call_usage AS usage
|
|
11031
|
+
ON usage.project_id = ${input.projectId}
|
|
11032
|
+
AND usage.execution_id = requested.execution_id
|
|
11033
|
+
`;
|
|
11034
|
+
}
|
|
11035
|
+
async assertExecutionExists(sql, projectId, executionId) {
|
|
11036
|
+
const [execution] = await sql`
|
|
11037
|
+
SELECT EXISTS (
|
|
11038
|
+
SELECT 1 FROM executions WHERE project_id = ${projectId} AND id = ${executionId}
|
|
11039
|
+
) AS exists
|
|
11040
|
+
`;
|
|
11041
|
+
if (execution?.exists)
|
|
11042
|
+
return;
|
|
11043
|
+
throw new AiUsageStorageError("missing_execution", `[SixbPg] AI usage execution '${executionId}' does not exist in project '${projectId}'.`);
|
|
11044
|
+
}
|
|
11045
|
+
}
|
|
11046
|
+
function aggregateExecutionRows(executionIds, rows) {
|
|
11047
|
+
const usageByExecutionId = new Map;
|
|
11048
|
+
for (const executionId of executionIds)
|
|
11049
|
+
usageByExecutionId.set(executionId, []);
|
|
11050
|
+
for (const row of rows) {
|
|
11051
|
+
usageByExecutionId.get(row.execution_id)?.push(usageFromRow(row));
|
|
11052
|
+
}
|
|
11053
|
+
return executionIds.map((executionId) => {
|
|
11054
|
+
const usages = usageByExecutionId.get(executionId) ?? [];
|
|
11055
|
+
return {
|
|
11056
|
+
modelCallCount: usages.length,
|
|
11057
|
+
usage: aggregateAiModelCallUsage(usages)
|
|
11058
|
+
};
|
|
11059
|
+
});
|
|
11060
|
+
}
|
|
11061
|
+
function usageFromRow(row) {
|
|
11062
|
+
return {
|
|
11063
|
+
...row.input_tokens === null ? {} : { inputTokens: Number(row.input_tokens) },
|
|
11064
|
+
...row.output_tokens === null ? {} : { outputTokens: Number(row.output_tokens) },
|
|
11065
|
+
...row.uncached_input_tokens === null ? {} : { uncachedInputTokens: Number(row.uncached_input_tokens) },
|
|
11066
|
+
...row.cache_read_input_tokens === null ? {} : { cacheReadInputTokens: Number(row.cache_read_input_tokens) },
|
|
11067
|
+
...row.cache_write_input_tokens === null ? {} : { cacheWriteInputTokens: Number(row.cache_write_input_tokens) },
|
|
11068
|
+
...row.text_output_tokens === null ? {} : { textOutputTokens: Number(row.text_output_tokens) },
|
|
11069
|
+
...row.reasoning_output_tokens === null ? {} : { reasoningOutputTokens: Number(row.reasoning_output_tokens) }
|
|
11070
|
+
};
|
|
11071
|
+
}
|
|
11072
|
+
function rawUsageFromRow(value) {
|
|
11073
|
+
return typeof value === "string" ? JSON.parse(value) : value;
|
|
11074
|
+
}
|
|
11075
|
+
function requestedReasoningFromRow(value) {
|
|
11076
|
+
const parsed = typeof value === "string" && (value.startsWith('"') || value.startsWith("{")) ? JSON.parse(value) : value;
|
|
11077
|
+
if (!isModelReasoning(parsed)) {
|
|
11078
|
+
throw new Error("[SixbPg] Stored AI usage requested reasoning is invalid.");
|
|
11079
|
+
}
|
|
11080
|
+
return parsed;
|
|
11081
|
+
}
|
|
11082
|
+
function assertNonBlankProjectId(projectId) {
|
|
11083
|
+
if (typeof projectId !== "string" || projectId.trim().length === 0) {
|
|
11084
|
+
throw new TypeError("[Sixb] AI usage projectId must be nonblank.");
|
|
11085
|
+
}
|
|
11086
|
+
}
|
|
11087
|
+
|
|
11088
|
+
// src/pg-ai-cost-storage.ts
|
|
11089
|
+
class PgAiCostStorage {
|
|
11090
|
+
sql;
|
|
11091
|
+
constructor(sql) {
|
|
11092
|
+
this.sql = sql;
|
|
11093
|
+
}
|
|
11094
|
+
async listModelCallGroups(input) {
|
|
11095
|
+
return listAiModelCallGroups(this.sql, input);
|
|
11096
|
+
}
|
|
11097
|
+
async recordModelCallCost(input) {
|
|
11098
|
+
const record = normalizeAiModelCallCostRecord(input);
|
|
11099
|
+
await runPgTransaction(this.sql, async (tx) => {
|
|
11100
|
+
const usage = await findUsage(tx, record.projectId, record.usageRecordId);
|
|
11101
|
+
if (!usage)
|
|
11102
|
+
throw missingUsage(record);
|
|
11103
|
+
assertCostMatchesUsage(record, usage);
|
|
11104
|
+
if (await findCost(tx, record.projectId, record.usageRecordId))
|
|
11105
|
+
return;
|
|
11106
|
+
const rated = record.status === "rated" ? record : undefined;
|
|
11107
|
+
const inserted = await tx`
|
|
11108
|
+
INSERT INTO ai_model_call_valuations (
|
|
11109
|
+
project_id, usage_record_id, status, provider_id, model_id,
|
|
10194
11110
|
currency, amount_nanos, reason, details, rated_at
|
|
10195
11111
|
) VALUES (
|
|
10196
11112
|
${record.projectId}, ${record.usageRecordId}, ${record.status},
|
|
@@ -10248,7 +11164,12 @@ class PgAiCostStorage {
|
|
|
10248
11164
|
usage.text_output_tokens,
|
|
10249
11165
|
usage.reasoning_output_tokens,
|
|
10250
11166
|
usage.reporting_status,
|
|
10251
|
-
|
|
11167
|
+
CASE
|
|
11168
|
+
WHEN direct_agent.kind = 'conversation' THEN 'agent'
|
|
11169
|
+
WHEN workflow_agent.node_run_id IS NOT NULL THEN 'workflowAgent'
|
|
11170
|
+
ELSE NULL
|
|
11171
|
+
END AS attribution_agent_kind,
|
|
11172
|
+
workflow_node.node_id AS attribution_agent_step_id,
|
|
10252
11173
|
workflow_node.workflow_id AS attribution_workflow_id,
|
|
10253
11174
|
cost.status AS cost_status,
|
|
10254
11175
|
cost.currency AS cost_currency,
|
|
@@ -10283,11 +11204,12 @@ class PgAiCostStorage {
|
|
|
10283
11204
|
SELECT 'totals'::text AS dimension,
|
|
10284
11205
|
NULL::text AS key_1,
|
|
10285
11206
|
NULL::text AS key_2,
|
|
11207
|
+
NULL::text AS key_3,
|
|
10286
11208
|
NULL::timestamptz AS bucket_start,
|
|
10287
11209
|
classified.*
|
|
10288
11210
|
FROM classified
|
|
10289
11211
|
UNION ALL
|
|
10290
|
-
SELECT 'series', NULL, NULL,
|
|
11212
|
+
SELECT 'series', NULL, NULL, NULL,
|
|
10291
11213
|
CASE ${input.bucket}::text
|
|
10292
11214
|
WHEN 'hour' THEN
|
|
10293
11215
|
date_trunc('hour', occurred_at AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'
|
|
@@ -10298,17 +11220,18 @@ class PgAiCostStorage {
|
|
|
10298
11220
|
classified.*
|
|
10299
11221
|
FROM classified
|
|
10300
11222
|
UNION ALL
|
|
10301
|
-
SELECT 'model', provider_id, requested_model_id, NULL, classified.* FROM classified
|
|
11223
|
+
SELECT 'model', provider_id, requested_model_id, NULL, NULL, classified.* FROM classified
|
|
10302
11224
|
UNION ALL
|
|
10303
|
-
SELECT 'agent',
|
|
10304
|
-
FROM classified WHERE
|
|
11225
|
+
SELECT 'agent', attribution_agent_kind, attribution_workflow_id, attribution_agent_step_id, NULL, classified.*
|
|
11226
|
+
FROM classified WHERE attribution_agent_kind IS NOT NULL
|
|
10305
11227
|
UNION ALL
|
|
10306
|
-
SELECT 'workflow', attribution_workflow_id, NULL, NULL, classified.*
|
|
11228
|
+
SELECT 'workflow', attribution_workflow_id, NULL, NULL, NULL, classified.*
|
|
10307
11229
|
FROM classified WHERE attribution_workflow_id IS NOT NULL
|
|
10308
11230
|
)
|
|
10309
11231
|
SELECT dimension,
|
|
10310
11232
|
key_1,
|
|
10311
11233
|
key_2,
|
|
11234
|
+
key_3,
|
|
10312
11235
|
bucket_start,
|
|
10313
11236
|
amount_currency,
|
|
10314
11237
|
COUNT(*)::text AS model_call_count,
|
|
@@ -10335,12 +11258,12 @@ class PgAiCostStorage {
|
|
|
10335
11258
|
AS unvalued_call_count,
|
|
10336
11259
|
SUM(amount_nanos)::text AS amount_nanos
|
|
10337
11260
|
FROM expanded
|
|
10338
|
-
GROUP BY dimension, key_1, key_2, bucket_start, amount_currency
|
|
11261
|
+
GROUP BY dimension, key_1, key_2, key_3, bucket_start, amount_currency
|
|
10339
11262
|
`;
|
|
10340
11263
|
return buildAiAccountingOverviewFromFragments(input, rows.map(aggregateFragmentFromRow));
|
|
10341
11264
|
}
|
|
10342
11265
|
async listModelCalls(input) {
|
|
10343
|
-
const query =
|
|
11266
|
+
const query = normalizeAiModelCallAccountingQuery2(input);
|
|
10344
11267
|
const providerId = query.providerId ?? null;
|
|
10345
11268
|
const modelId = query.modelId ?? null;
|
|
10346
11269
|
const executionId = query.executionId ?? null;
|
|
@@ -10366,13 +11289,15 @@ class PgAiCostStorage {
|
|
|
10366
11289
|
const rows = await this.sql`
|
|
10367
11290
|
SELECT usage.*,
|
|
10368
11291
|
CASE
|
|
10369
|
-
WHEN direct_agent.
|
|
11292
|
+
WHEN direct_agent.kind = 'conversation' THEN 'agent'
|
|
11293
|
+
WHEN direct_agent.kind = 'subagent' THEN 'subagent'
|
|
10370
11294
|
WHEN workflow_agent.node_run_id IS NOT NULL THEN 'workflowAgent'
|
|
10371
11295
|
ELSE NULL
|
|
10372
11296
|
END AS attribution_kind,
|
|
10373
|
-
|
|
11297
|
+
workflow_node.node_id AS attribution_agent_step_id,
|
|
10374
11298
|
direct_agent.id AS attribution_agent_run_id,
|
|
10375
11299
|
direct_agent.thread_id AS attribution_thread_id,
|
|
11300
|
+
direct_agent.parent_run_id AS attribution_parent_run_id,
|
|
10376
11301
|
workflow_agent.node_run_id AS attribution_node_run_id,
|
|
10377
11302
|
workflow_node.workflow_id AS attribution_workflow_id,
|
|
10378
11303
|
workflow_node.workflow_run_id AS attribution_workflow_run_id,
|
|
@@ -10438,10 +11363,13 @@ async function findUsage(sql, projectId, id) {
|
|
|
10438
11363
|
const [row] = await sql`
|
|
10439
11364
|
SELECT * FROM ai_model_call_usage WHERE project_id = ${projectId} AND id = ${id}
|
|
10440
11365
|
`;
|
|
10441
|
-
return row ?
|
|
11366
|
+
return row ? usageFromRow2(row, []) : null;
|
|
10442
11367
|
}
|
|
10443
11368
|
function aggregateFragmentFromRow(row) {
|
|
10444
|
-
const dimensionKeys = row.dimension === "model" ? { providerId: required2(row.key_1), modelId: required2(row.key_2) } : row.dimension === "agent" ? {
|
|
11369
|
+
const dimensionKeys = row.dimension === "model" ? { providerId: required2(row.key_1), modelId: required2(row.key_2) } : row.dimension === "agent" ? {
|
|
11370
|
+
agentKind: required2(row.key_1),
|
|
11371
|
+
...row.key_1 === "workflowAgent" ? { workflowId: required2(row.key_2), agentStepId: required2(row.key_3) } : {}
|
|
11372
|
+
} : row.dimension === "workflow" ? { workflowId: required2(row.key_1) } : row.dimension === "series" ? { start: new Date(required2(row.bucket_start)) } : {};
|
|
10445
11373
|
return {
|
|
10446
11374
|
dimension: row.dimension,
|
|
10447
11375
|
...dimensionKeys,
|
|
@@ -10473,16 +11401,19 @@ function meterFragment(presentCallCount, total) {
|
|
|
10473
11401
|
return { presentCallCount, total };
|
|
10474
11402
|
}
|
|
10475
11403
|
function accountingItemFromRow(row) {
|
|
10476
|
-
const usage =
|
|
11404
|
+
const usage = usageFromRow2(row, parseStringArray(row.requester_group_ids));
|
|
10477
11405
|
const cost = row.cost_status === null ? undefined : costFromJoinedRow(row);
|
|
10478
11406
|
const attribution = row.attribution_kind === null ? undefined : row.attribution_kind === "agent" ? {
|
|
10479
11407
|
kind: "agent",
|
|
10480
|
-
agentId: required2(row.attribution_agent_id),
|
|
10481
11408
|
agentRunId: required2(row.attribution_agent_run_id),
|
|
10482
11409
|
threadId: required2(row.attribution_thread_id)
|
|
11410
|
+
} : row.attribution_kind === "subagent" ? {
|
|
11411
|
+
kind: "subagent",
|
|
11412
|
+
subagentRunId: required2(row.attribution_agent_run_id),
|
|
11413
|
+
parentRunId: required2(row.attribution_parent_run_id)
|
|
10483
11414
|
} : {
|
|
10484
11415
|
kind: "workflowAgent",
|
|
10485
|
-
|
|
11416
|
+
agentStepId: required2(row.attribution_agent_step_id),
|
|
10486
11417
|
nodeRunId: required2(row.attribution_node_run_id),
|
|
10487
11418
|
workflowId: required2(row.attribution_workflow_id),
|
|
10488
11419
|
workflowRunId: required2(row.attribution_workflow_run_id)
|
|
@@ -10493,7 +11424,7 @@ function accountingItemFromRow(row) {
|
|
|
10493
11424
|
...cost ? { cost } : {}
|
|
10494
11425
|
};
|
|
10495
11426
|
}
|
|
10496
|
-
function
|
|
11427
|
+
function usageFromRow2(row, requesterGroupIds) {
|
|
10497
11428
|
const number = (value) => value === null ? undefined : Number(value);
|
|
10498
11429
|
return {
|
|
10499
11430
|
id: row.id,
|
|
@@ -10503,7 +11434,11 @@ function usageFromRow(row, requesterGroupIds) {
|
|
|
10503
11434
|
callId: row.call_id,
|
|
10504
11435
|
requesterGroupIds,
|
|
10505
11436
|
providerId: row.provider_id,
|
|
11437
|
+
...row.provider_ids === null ? {} : {
|
|
11438
|
+
providerIds: normalizeModelProviderIds2(typeof row.provider_ids === "string" ? JSON.parse(row.provider_ids) : row.provider_ids)
|
|
11439
|
+
},
|
|
10506
11440
|
requestedModelId: row.requested_model_id,
|
|
11441
|
+
...row.requested_reasoning === null ? {} : { requestedReasoning: requestedReasoningFromRow(row.requested_reasoning) },
|
|
10507
11442
|
...row.response_model_id === null ? {} : { responseModelId: row.response_model_id },
|
|
10508
11443
|
responseId: row.response_id,
|
|
10509
11444
|
usage: {
|
|
@@ -10552,14 +11487,12 @@ function costFromJoinedRow(row) {
|
|
|
10552
11487
|
}
|
|
10553
11488
|
function costRecord(input) {
|
|
10554
11489
|
const details = parseAiModelCallCostDetails(input.details);
|
|
11490
|
+
const priceSource = details.priceSource === undefined ? undefined : { ...details.priceSource, observedAt: new Date(details.priceSource.observedAt) };
|
|
10555
11491
|
const base = {
|
|
11492
|
+
...details.estimate === undefined ? {} : { estimate: details.estimate },
|
|
10556
11493
|
projectId: input.projectId,
|
|
10557
11494
|
usageRecordId: input.usageId,
|
|
10558
11495
|
pricingContext: details.pricingContext,
|
|
10559
|
-
priceSource: {
|
|
10560
|
-
...details.priceSource,
|
|
10561
|
-
observedAt: new Date(details.priceSource.observedAt)
|
|
10562
|
-
},
|
|
10563
11496
|
ratedAt: new Date(input.ratedAt)
|
|
10564
11497
|
};
|
|
10565
11498
|
if (input.status === "rated") {
|
|
@@ -10567,6 +11500,7 @@ function costRecord(input) {
|
|
|
10567
11500
|
...base,
|
|
10568
11501
|
status: "rated",
|
|
10569
11502
|
billingIdentity: { providerId: required2(input.providerId), modelId: required2(input.modelId) },
|
|
11503
|
+
priceSource: required2(priceSource),
|
|
10570
11504
|
money: { currency: required2(input.currency), amountNanos: required2(input.amountNanos) },
|
|
10571
11505
|
components: required2(details.components)
|
|
10572
11506
|
});
|
|
@@ -10575,6 +11509,7 @@ function costRecord(input) {
|
|
|
10575
11509
|
...base,
|
|
10576
11510
|
status: "unpriceable",
|
|
10577
11511
|
...input.providerId && input.modelId ? { billingIdentity: { providerId: input.providerId, modelId: input.modelId } } : {},
|
|
11512
|
+
...priceSource === undefined ? {} : { priceSource },
|
|
10578
11513
|
reason: required2(input.reason),
|
|
10579
11514
|
...details.missingMeters ? { missingMeters: details.missingMeters } : {}
|
|
10580
11515
|
});
|
|
@@ -10655,222 +11590,611 @@ function required2(value) {
|
|
|
10655
11590
|
return value;
|
|
10656
11591
|
}
|
|
10657
11592
|
|
|
10658
|
-
// src/pg-ai-
|
|
11593
|
+
// src/pg-ai-limit-storage.ts
|
|
10659
11594
|
import {
|
|
10660
|
-
|
|
10661
|
-
|
|
10662
|
-
|
|
10663
|
-
|
|
10664
|
-
|
|
10665
|
-
|
|
11595
|
+
AiLimitOperationLock,
|
|
11596
|
+
aiLimitAmountKey,
|
|
11597
|
+
aiLimitQuantityFromAmount,
|
|
11598
|
+
aiLimitReservationBuckets,
|
|
11599
|
+
aiLimitReservationRequestKey,
|
|
11600
|
+
aiLimitReservationRequestMatches,
|
|
11601
|
+
aiLimitSubjectFromParts,
|
|
11602
|
+
aiLimitSubjectKey,
|
|
11603
|
+
aiLimitSubjectParts,
|
|
11604
|
+
aiLimitSubjectsFromAccountingEntry,
|
|
11605
|
+
assertAiLimitAccountingMatchesReservation,
|
|
11606
|
+
assertNonBlank as assertNonBlank2,
|
|
11607
|
+
cloneValidDate,
|
|
11608
|
+
normalizeAiLimitAmount,
|
|
11609
|
+
normalizeAiLimitQuantities,
|
|
11610
|
+
normalizeAiModelCallReservationIdentity,
|
|
11611
|
+
normalizeCreateAiLimitPolicy,
|
|
11612
|
+
normalizeReserveAiModelCall,
|
|
11613
|
+
normalizeUpdateAiLimitPolicy,
|
|
11614
|
+
resolveAiLimitActual
|
|
11615
|
+
} from "@sixb/core/internal/ai-limit-storage-provider";
|
|
11616
|
+
import { AiLimitStorageError, aiLimitCalendarMonth } from "@sixb/core/storage";
|
|
11617
|
+
class PgAiLimitStorage {
|
|
10666
11618
|
sql;
|
|
11619
|
+
operations = new AiLimitOperationLock;
|
|
10667
11620
|
constructor(sql) {
|
|
10668
11621
|
this.sql = sql;
|
|
10669
11622
|
}
|
|
10670
|
-
|
|
10671
|
-
|
|
10672
|
-
|
|
10673
|
-
|
|
10674
|
-
|
|
10675
|
-
|
|
10676
|
-
|
|
11623
|
+
settleOperations() {
|
|
11624
|
+
return this.operations.settle();
|
|
11625
|
+
}
|
|
11626
|
+
async createPolicy(input) {
|
|
11627
|
+
const policy = normalizeCreateAiLimitPolicy(input);
|
|
11628
|
+
return this.withProjectLock(policy.projectId, async (tx) => {
|
|
11629
|
+
const subject = aiLimitSubjectParts(policy.subject);
|
|
11630
|
+
const limit = normalizeAiLimitAmount(policy.limit);
|
|
10677
11631
|
const inserted = await tx`
|
|
10678
|
-
INSERT INTO
|
|
10679
|
-
project_id,
|
|
10680
|
-
|
|
10681
|
-
execution_id,
|
|
10682
|
-
attempt,
|
|
10683
|
-
call_id,
|
|
10684
|
-
provider_id,
|
|
10685
|
-
requested_model_id,
|
|
10686
|
-
response_model_id,
|
|
10687
|
-
response_id,
|
|
10688
|
-
input_tokens,
|
|
10689
|
-
output_tokens,
|
|
10690
|
-
total_tokens,
|
|
10691
|
-
uncached_input_tokens,
|
|
10692
|
-
cache_read_input_tokens,
|
|
10693
|
-
cache_write_input_tokens,
|
|
10694
|
-
text_output_tokens,
|
|
10695
|
-
reasoning_output_tokens,
|
|
10696
|
-
reporting_status,
|
|
10697
|
-
raw_usage,
|
|
10698
|
-
occurred_at,
|
|
10699
|
-
recorded_at
|
|
11632
|
+
INSERT INTO ai_usage_limits (
|
|
11633
|
+
project_id, id, subject_type, subject_id, meter, currency, limit_amount,
|
|
11634
|
+
period_kind, enabled, created_at, updated_at
|
|
10700
11635
|
) VALUES (
|
|
10701
|
-
${
|
|
10702
|
-
${
|
|
10703
|
-
${
|
|
10704
|
-
${record.attempt},
|
|
10705
|
-
${record.callId},
|
|
10706
|
-
${record.providerId},
|
|
10707
|
-
${record.requestedModelId},
|
|
10708
|
-
${record.responseModelId ?? null},
|
|
10709
|
-
${record.responseId},
|
|
10710
|
-
${record.usage.inputTokens ?? null},
|
|
10711
|
-
${record.usage.outputTokens ?? null},
|
|
10712
|
-
${record.usage.totalTokens ?? null},
|
|
10713
|
-
${record.usage.uncachedInputTokens ?? null},
|
|
10714
|
-
${record.usage.cacheReadInputTokens ?? null},
|
|
10715
|
-
${record.usage.cacheWriteInputTokens ?? null},
|
|
10716
|
-
${record.usage.textOutputTokens ?? null},
|
|
10717
|
-
${record.usage.reasoningOutputTokens ?? null},
|
|
10718
|
-
${record.usage.reportingStatus},
|
|
10719
|
-
${record.rawUsage === undefined ? null : JSON.stringify(record.rawUsage)}::text::jsonb,
|
|
10720
|
-
${record.occurredAt},
|
|
10721
|
-
${record.recordedAt}
|
|
11636
|
+
${policy.projectId}, ${policy.id}, ${subject.type}, ${subject.id}, ${limit.meter},
|
|
11637
|
+
${limit.currency}, ${limit.amount.toString()}, ${policy.period}, ${policy.enabled},
|
|
11638
|
+
${policy.createdAt}, ${policy.updatedAt}
|
|
10722
11639
|
)
|
|
10723
11640
|
ON CONFLICT DO NOTHING
|
|
10724
11641
|
RETURNING id
|
|
10725
11642
|
`;
|
|
10726
|
-
if (inserted.length === 0)
|
|
10727
|
-
|
|
10728
|
-
|
|
10729
|
-
|
|
10730
|
-
|
|
11643
|
+
if (inserted.length === 0)
|
|
11644
|
+
throw duplicatePolicy(policy);
|
|
11645
|
+
return this.requirePolicy(tx, policy.projectId, policy.id);
|
|
11646
|
+
});
|
|
11647
|
+
}
|
|
11648
|
+
async updatePolicy(input) {
|
|
11649
|
+
const update = normalizeUpdateAiLimitPolicy(input);
|
|
11650
|
+
return this.withProjectLock(update.projectId, async (tx) => {
|
|
11651
|
+
const existing = await this.requirePolicy(tx, update.projectId, update.id);
|
|
11652
|
+
if (update.limit !== undefined)
|
|
11653
|
+
assertSameDimension(existing.limit, update.limit);
|
|
11654
|
+
const limit = normalizeAiLimitAmount(update.limit ?? existing.limit);
|
|
11655
|
+
await tx`
|
|
11656
|
+
UPDATE ai_usage_limits
|
|
11657
|
+
SET limit_amount = ${limit.amount.toString()},
|
|
11658
|
+
enabled = ${update.enabled ?? existing.enabled},
|
|
11659
|
+
updated_at = ${update.updatedAt}
|
|
11660
|
+
WHERE project_id = ${update.projectId} AND id = ${update.id}
|
|
11661
|
+
`;
|
|
11662
|
+
return this.requirePolicy(tx, update.projectId, update.id);
|
|
11663
|
+
});
|
|
11664
|
+
}
|
|
11665
|
+
async deletePolicy(input) {
|
|
11666
|
+
validatePolicyIdentity(input);
|
|
11667
|
+
return this.withProjectLock(input.projectId, async (tx) => {
|
|
11668
|
+
const deleted = await tx`
|
|
11669
|
+
DELETE FROM ai_usage_limits
|
|
11670
|
+
WHERE project_id = ${input.projectId} AND id = ${input.id}
|
|
11671
|
+
RETURNING id
|
|
11672
|
+
`;
|
|
11673
|
+
return deleted.length > 0;
|
|
11674
|
+
});
|
|
11675
|
+
}
|
|
11676
|
+
async getPolicy(input) {
|
|
11677
|
+
validatePolicyIdentity(input);
|
|
11678
|
+
const [row] = await this.sql`
|
|
11679
|
+
SELECT * FROM ai_usage_limits
|
|
11680
|
+
WHERE project_id = ${input.projectId} AND id = ${input.id}
|
|
11681
|
+
`;
|
|
11682
|
+
return row ? policyFromRow(row) : null;
|
|
11683
|
+
}
|
|
11684
|
+
async listPolicies(input) {
|
|
11685
|
+
assertNonBlank2(input.projectId, "projectId");
|
|
11686
|
+
return (await this.policyRows(this.sql, input.projectId, input.includeDisabled ?? false)).map(policyFromRow);
|
|
11687
|
+
}
|
|
11688
|
+
async listPolicyStatuses(input) {
|
|
11689
|
+
assertNonBlank2(input.projectId, "projectId");
|
|
11690
|
+
const period = aiLimitCalendarMonth(input.at ?? new Date);
|
|
11691
|
+
const existingGroups = input.existingGroupIds === undefined ? undefined : new Set(input.existingGroupIds);
|
|
11692
|
+
return this.withProjectLock(input.projectId, async (tx) => {
|
|
11693
|
+
const rows = await this.policyRows(tx, input.projectId, input.includeDisabled ?? false);
|
|
11694
|
+
return Promise.all(rows.map((row) => this.statusForPolicy(tx, policyFromRow(row), period, existingGroups)));
|
|
11695
|
+
});
|
|
11696
|
+
}
|
|
11697
|
+
async reserveModelCall(input) {
|
|
11698
|
+
const request = normalizeReserveAiModelCall(input);
|
|
11699
|
+
return this.withProjectLock(request.identity.projectId, async (tx) => {
|
|
11700
|
+
const existing = await this.findReservation(tx, request.identity);
|
|
11701
|
+
if (existing) {
|
|
11702
|
+
if (!aiLimitReservationRequestMatches(existing, request)) {
|
|
11703
|
+
throw new AiLimitStorageError("reservation_conflict", `[SixbPg] AI model-call reservation '${input.callId}' was replayed with different subjects, estimates, or period.`);
|
|
11704
|
+
}
|
|
11705
|
+
if (existing.state !== "active") {
|
|
11706
|
+
return { status: "terminal", reservation: existing, created: false };
|
|
11707
|
+
}
|
|
11708
|
+
return { status: "reserved", reservation: existing, created: false };
|
|
10731
11709
|
}
|
|
10732
|
-
|
|
10733
|
-
|
|
10734
|
-
|
|
10735
|
-
|
|
10736
|
-
|
|
10737
|
-
|
|
10738
|
-
|
|
11710
|
+
await this.assertExecutionExists(tx, request.identity.projectId, request.identity.executionId);
|
|
11711
|
+
const subjectKeys = new Set(request.subjects.map(aiLimitSubjectKey));
|
|
11712
|
+
const estimates = new Map(request.estimates.map((quantity) => {
|
|
11713
|
+
const amount = normalizeAiLimitAmount(quantity);
|
|
11714
|
+
return [aiLimitAmountKey(amount), amount];
|
|
11715
|
+
}));
|
|
11716
|
+
const enabledPolicies = (await this.policyRows(tx, request.identity.projectId, false)).map(policyFromRow);
|
|
11717
|
+
const exhaustedPolicies = [];
|
|
11718
|
+
const unavailablePolicies = [];
|
|
11719
|
+
const unavailableReasons = new Set;
|
|
11720
|
+
for (const policy of enabledPolicies) {
|
|
11721
|
+
if (!subjectKeys.has(aiLimitSubjectKey(policy.subject)))
|
|
11722
|
+
continue;
|
|
11723
|
+
const limit = normalizeAiLimitAmount(policy.limit);
|
|
11724
|
+
const estimate = estimates.get(aiLimitAmountKey(limit));
|
|
11725
|
+
const status = await this.statusForPolicy(tx, policy, request.period);
|
|
11726
|
+
if (!estimate || status.accountingStatus === "unavailable") {
|
|
11727
|
+
unavailablePolicies.push(status);
|
|
11728
|
+
if (!estimate)
|
|
11729
|
+
unavailableReasons.add("missingEstimate");
|
|
11730
|
+
if (status.accountingStatus === "unavailable") {
|
|
11731
|
+
unavailableReasons.add("incompleteAccounting");
|
|
11732
|
+
}
|
|
11733
|
+
continue;
|
|
11734
|
+
}
|
|
11735
|
+
const consumed = normalizeAiLimitAmount(status.consumption.actual).amount + normalizeAiLimitAmount(status.consumption.reserved).amount + normalizeAiLimitAmount(status.consumption.unknown).amount;
|
|
11736
|
+
if (consumed + estimate.amount > limit.amount)
|
|
11737
|
+
exhaustedPolicies.push(status);
|
|
11738
|
+
}
|
|
11739
|
+
if (unavailablePolicies.length > 0) {
|
|
11740
|
+
unavailablePolicies.sort((left, right) => left.policy.id.localeCompare(right.policy.id));
|
|
11741
|
+
return {
|
|
11742
|
+
status: "unavailable",
|
|
11743
|
+
unavailablePolicies,
|
|
11744
|
+
reasons: [...unavailableReasons].sort()
|
|
11745
|
+
};
|
|
11746
|
+
}
|
|
11747
|
+
if (exhaustedPolicies.length > 0) {
|
|
11748
|
+
exhaustedPolicies.sort((left, right) => left.policy.id.localeCompare(right.policy.id));
|
|
11749
|
+
return {
|
|
11750
|
+
status: "denied",
|
|
11751
|
+
exhaustedPolicies,
|
|
11752
|
+
resetAt: new Date(Math.min(...exhaustedPolicies.map((status) => status.period.resetAt.getTime())))
|
|
11753
|
+
};
|
|
11754
|
+
}
|
|
11755
|
+
const buckets = aiLimitReservationBuckets(enabledPolicies, request.subjects, request.estimates);
|
|
11756
|
+
if (buckets.length === 0)
|
|
11757
|
+
return { status: "notRequired" };
|
|
11758
|
+
for (const bucket of buckets) {
|
|
11759
|
+
const amount = normalizeAiLimitAmount(bucket.estimate);
|
|
11760
|
+
await this.changeCounters(tx, request.identity.projectId, bucket.subject, amount, request.period, { reserved: amount.amount }, request.reservedAt);
|
|
11761
|
+
}
|
|
11762
|
+
const reservation = {
|
|
11763
|
+
...request.identity,
|
|
11764
|
+
buckets,
|
|
11765
|
+
requestKey: aiLimitReservationRequestKey(request),
|
|
11766
|
+
period: request.period,
|
|
11767
|
+
state: "active",
|
|
11768
|
+
reservedAt: request.reservedAt,
|
|
11769
|
+
updatedAt: new Date(request.reservedAt)
|
|
11770
|
+
};
|
|
11771
|
+
await this.insertReservation(tx, reservation);
|
|
11772
|
+
return { status: "reserved", reservation: structuredClone(reservation), created: true };
|
|
11773
|
+
});
|
|
11774
|
+
}
|
|
11775
|
+
async recordModelCallActuals(input) {
|
|
11776
|
+
assertNonBlank2(input.projectId, "projectId");
|
|
11777
|
+
assertNonBlank2(input.usageRecordId, "usageRecordId");
|
|
11778
|
+
const recordedAt = cloneValidDate(input.recordedAt ?? new Date, "recordedAt");
|
|
11779
|
+
await this.withProjectLock(input.projectId, async (tx) => {
|
|
11780
|
+
const entry = await this.requireAccountingEntry(tx, input.projectId, input.usageRecordId);
|
|
11781
|
+
const period = aiLimitCalendarMonth(entry.occurredAt);
|
|
11782
|
+
for (const subjectValue of aiLimitSubjectsFromAccountingEntry(entry)) {
|
|
11783
|
+
const subject = aiLimitSubjectParts(subjectValue);
|
|
11784
|
+
const rows = await tx`
|
|
11785
|
+
SELECT meter, currency, actual_amount, accounting_status
|
|
11786
|
+
FROM ai_usage_limit_periods
|
|
11787
|
+
WHERE project_id = ${input.projectId}
|
|
11788
|
+
AND subject_type = ${subject.type}
|
|
11789
|
+
AND subject_id = ${subject.id}
|
|
11790
|
+
AND period_start = ${period.start}
|
|
11791
|
+
FOR UPDATE
|
|
10739
11792
|
`;
|
|
11793
|
+
for (const row of rows) {
|
|
11794
|
+
const actual = resolveAiLimitActual([entry], row);
|
|
11795
|
+
await tx`
|
|
11796
|
+
UPDATE ai_usage_limit_periods
|
|
11797
|
+
SET actual_amount = ${(BigInt(row.actual_amount) + actual.amount).toString()},
|
|
11798
|
+
accounting_status = ${row.accounting_status === "unavailable" || actual.accountingStatus === "unavailable" ? "unavailable" : "complete"},
|
|
11799
|
+
updated_at = ${recordedAt}
|
|
11800
|
+
WHERE project_id = ${input.projectId}
|
|
11801
|
+
AND subject_type = ${subject.type}
|
|
11802
|
+
AND subject_id = ${subject.id}
|
|
11803
|
+
AND meter = ${row.meter}
|
|
11804
|
+
AND currency = ${row.currency}
|
|
11805
|
+
AND period_start = ${period.start}
|
|
11806
|
+
`;
|
|
11807
|
+
}
|
|
10740
11808
|
}
|
|
10741
|
-
return { record: await this.requireById(tx, record.projectId, record.id), created: true };
|
|
10742
11809
|
});
|
|
10743
11810
|
}
|
|
10744
|
-
async
|
|
10745
|
-
const
|
|
10746
|
-
|
|
10747
|
-
|
|
11811
|
+
async reconcileModelCall(input) {
|
|
11812
|
+
const identity = normalizeAiModelCallReservationIdentity(input);
|
|
11813
|
+
assertNonBlank2(input.usageRecordId, "usageRecordId");
|
|
11814
|
+
const reconciledAt = cloneValidDate(input.reconciledAt ?? new Date, "reconciledAt");
|
|
11815
|
+
return this.withProjectLock(identity.projectId, async (tx) => {
|
|
11816
|
+
const reservation = await this.findReservation(tx, identity);
|
|
11817
|
+
if (!reservation)
|
|
11818
|
+
throw missingReservation(identity);
|
|
11819
|
+
if (reservation.state === "reconciled") {
|
|
11820
|
+
if (reservation.usageRecordId === input.usageRecordId) {
|
|
11821
|
+
return reservation;
|
|
11822
|
+
}
|
|
11823
|
+
throw reconciliationConflict(identity.callId);
|
|
11824
|
+
}
|
|
11825
|
+
const accounting = await this.requireAccountingEntry(tx, identity.projectId, input.usageRecordId);
|
|
11826
|
+
assertAiLimitAccountingMatchesReservation(reservation, accounting);
|
|
11827
|
+
await this.reconcileEstimateCounters(tx, reservation, accounting, reconciledAt);
|
|
11828
|
+
const reconciled = {
|
|
11829
|
+
...reservation,
|
|
11830
|
+
state: "reconciled",
|
|
11831
|
+
usageRecordId: input.usageRecordId,
|
|
11832
|
+
updatedAt: reconciledAt
|
|
11833
|
+
};
|
|
11834
|
+
await this.updateReservation(tx, reconciled);
|
|
11835
|
+
return structuredClone(reconciled);
|
|
10748
11836
|
});
|
|
10749
|
-
return summary;
|
|
10750
11837
|
}
|
|
10751
|
-
async
|
|
10752
|
-
|
|
10753
|
-
|
|
10754
|
-
|
|
10755
|
-
|
|
10756
|
-
|
|
10757
|
-
|
|
10758
|
-
|
|
10759
|
-
|
|
10760
|
-
|
|
11838
|
+
async markReservationUnknown(input) {
|
|
11839
|
+
const identity = normalizeAiModelCallReservationIdentity(input);
|
|
11840
|
+
const markedAt = cloneValidDate(input.markedAt ?? new Date, "markedAt");
|
|
11841
|
+
return this.withProjectLock(identity.projectId, async (tx) => {
|
|
11842
|
+
const reservation = await this.findReservation(tx, identity);
|
|
11843
|
+
if (!reservation)
|
|
11844
|
+
throw missingReservation(identity);
|
|
11845
|
+
if (reservation.state === "unknown")
|
|
11846
|
+
return reservation;
|
|
11847
|
+
if (reservation.state !== "active")
|
|
11848
|
+
throw invalidState(reservation, "mark unknown");
|
|
11849
|
+
await this.moveEstimateCounters(tx, reservation, { reserved: -1n, unknown: 1n }, markedAt);
|
|
11850
|
+
const unknown = { ...reservation, state: "unknown", updatedAt: markedAt };
|
|
11851
|
+
await this.updateReservation(tx, unknown);
|
|
11852
|
+
return structuredClone(unknown);
|
|
11853
|
+
});
|
|
10761
11854
|
}
|
|
10762
|
-
async
|
|
10763
|
-
|
|
10764
|
-
|
|
10765
|
-
|
|
10766
|
-
|
|
10767
|
-
|
|
10768
|
-
|
|
10769
|
-
|
|
11855
|
+
async withProjectLock(projectId, run) {
|
|
11856
|
+
return runPgTransaction(this.sql, async (tx) => {
|
|
11857
|
+
const operation = async () => {
|
|
11858
|
+
await lockAdvisoryKeys(tx, [`ai-limit:${projectId}`]);
|
|
11859
|
+
return run(tx);
|
|
11860
|
+
};
|
|
11861
|
+
return tx === this.sql ? this.operations.run(operation) : operation();
|
|
11862
|
+
});
|
|
10770
11863
|
}
|
|
10771
|
-
|
|
10772
|
-
|
|
10773
|
-
SELECT * FROM
|
|
10774
|
-
WHERE project_id = ${
|
|
10775
|
-
|
|
10776
|
-
AND attempt = ${record.attempt}
|
|
10777
|
-
AND call_id = ${record.callId}
|
|
10778
|
-
AND response_id = ${record.responseId}
|
|
11864
|
+
policyRows(sql, projectId, includeDisabled) {
|
|
11865
|
+
return sql`
|
|
11866
|
+
SELECT * FROM ai_usage_limits
|
|
11867
|
+
WHERE project_id = ${projectId} AND (${includeDisabled} OR enabled)
|
|
11868
|
+
ORDER BY id
|
|
10779
11869
|
`;
|
|
10780
|
-
return rows[0] ? this.rowToRecord(sql, rows[0]) : null;
|
|
10781
11870
|
}
|
|
10782
|
-
async
|
|
11871
|
+
async requirePolicy(sql, projectId, id) {
|
|
10783
11872
|
const [row] = await sql`
|
|
10784
|
-
SELECT * FROM
|
|
11873
|
+
SELECT * FROM ai_usage_limits WHERE project_id = ${projectId} AND id = ${id}
|
|
11874
|
+
`;
|
|
11875
|
+
if (row)
|
|
11876
|
+
return policyFromRow(row);
|
|
11877
|
+
throw new AiLimitStorageError("missing_policy", `[SixbPg] AI limit policy '${id}' does not exist in project '${projectId}'.`);
|
|
11878
|
+
}
|
|
11879
|
+
async statusForPolicy(sql, policy, period, existingGroups) {
|
|
11880
|
+
const subject = aiLimitSubjectParts(policy.subject);
|
|
11881
|
+
const limit = normalizeAiLimitAmount(policy.limit);
|
|
11882
|
+
let [row] = await sql`
|
|
11883
|
+
SELECT actual_amount, reserved_amount, unknown_amount, accounting_status
|
|
11884
|
+
FROM ai_usage_limit_periods
|
|
11885
|
+
WHERE project_id = ${policy.projectId}
|
|
11886
|
+
AND subject_type = ${subject.type}
|
|
11887
|
+
AND subject_id = ${subject.id}
|
|
11888
|
+
AND meter = ${limit.meter}
|
|
11889
|
+
AND currency = ${limit.currency}
|
|
11890
|
+
AND period_start = ${period.start}
|
|
10785
11891
|
`;
|
|
11892
|
+
if (!row || row.accounting_status === "unavailable") {
|
|
11893
|
+
const actual = resolveAiLimitActual(await this.accountingEntries(sql, policy, period), limit);
|
|
11894
|
+
const [inserted] = await sql`
|
|
11895
|
+
INSERT INTO ai_usage_limit_periods (
|
|
11896
|
+
project_id, subject_type, subject_id, meter, currency, period_start, period_end,
|
|
11897
|
+
actual_amount, reserved_amount, unknown_amount, accounting_status, updated_at
|
|
11898
|
+
) VALUES (
|
|
11899
|
+
${policy.projectId}, ${subject.type}, ${subject.id}, ${limit.meter}, ${limit.currency},
|
|
11900
|
+
${period.start}, ${period.end}, ${actual.amount.toString()}, 0, 0,
|
|
11901
|
+
${actual.accountingStatus}, ${new Date}
|
|
11902
|
+
)
|
|
11903
|
+
ON CONFLICT (project_id, subject_type, subject_id, meter, currency, period_start)
|
|
11904
|
+
DO UPDATE SET
|
|
11905
|
+
actual_amount = excluded.actual_amount,
|
|
11906
|
+
accounting_status = excluded.accounting_status,
|
|
11907
|
+
updated_at = excluded.updated_at
|
|
11908
|
+
RETURNING actual_amount, reserved_amount, unknown_amount, accounting_status
|
|
11909
|
+
`;
|
|
11910
|
+
row = inserted;
|
|
11911
|
+
}
|
|
10786
11912
|
if (!row)
|
|
10787
|
-
throw new Error(
|
|
10788
|
-
return
|
|
11913
|
+
throw new Error("[SixbPg] Failed to initialize AI limit period state.");
|
|
11914
|
+
return policyStatus(policy, period, row, existingGroups);
|
|
10789
11915
|
}
|
|
10790
|
-
async
|
|
10791
|
-
const
|
|
10792
|
-
|
|
10793
|
-
|
|
10794
|
-
|
|
11916
|
+
async changeCounters(sql, projectId, subjectValue, amount, period, changes, updatedAt) {
|
|
11917
|
+
const subject = aiLimitSubjectParts(subjectValue);
|
|
11918
|
+
const [row] = await sql`
|
|
11919
|
+
SELECT actual_amount, reserved_amount, unknown_amount, accounting_status
|
|
11920
|
+
FROM ai_usage_limit_periods
|
|
11921
|
+
WHERE project_id = ${projectId}
|
|
11922
|
+
AND subject_type = ${subject.type}
|
|
11923
|
+
AND subject_id = ${subject.id}
|
|
11924
|
+
AND meter = ${amount.meter}
|
|
11925
|
+
AND currency = ${amount.currency}
|
|
11926
|
+
AND period_start = ${period.start}
|
|
11927
|
+
FOR UPDATE
|
|
11928
|
+
`;
|
|
11929
|
+
const reserved = BigInt(row?.reserved_amount ?? "0") + (changes.reserved ?? 0n);
|
|
11930
|
+
const unknown = BigInt(row?.unknown_amount ?? "0") + (changes.unknown ?? 0n);
|
|
11931
|
+
if (reserved < 0n || unknown < 0n) {
|
|
11932
|
+
throw new Error("[SixbPg] AI limit period counter became negative.");
|
|
11933
|
+
}
|
|
11934
|
+
await sql`
|
|
11935
|
+
INSERT INTO ai_usage_limit_periods (
|
|
11936
|
+
project_id, subject_type, subject_id, meter, currency, period_start, period_end,
|
|
11937
|
+
actual_amount, reserved_amount, unknown_amount, accounting_status, updated_at
|
|
11938
|
+
) VALUES (
|
|
11939
|
+
${projectId}, ${subject.type}, ${subject.id}, ${amount.meter}, ${amount.currency},
|
|
11940
|
+
${period.start}, ${period.end}, ${row?.actual_amount ?? "0"}, ${reserved.toString()},
|
|
11941
|
+
${unknown.toString()}, ${row?.accounting_status ?? "complete"}, ${updatedAt}
|
|
11942
|
+
)
|
|
11943
|
+
ON CONFLICT (project_id, subject_type, subject_id, meter, currency, period_start)
|
|
11944
|
+
DO UPDATE SET
|
|
11945
|
+
period_end = excluded.period_end,
|
|
11946
|
+
actual_amount = excluded.actual_amount,
|
|
11947
|
+
reserved_amount = excluded.reserved_amount,
|
|
11948
|
+
unknown_amount = excluded.unknown_amount,
|
|
11949
|
+
accounting_status = excluded.accounting_status,
|
|
11950
|
+
updated_at = excluded.updated_at
|
|
10795
11951
|
`;
|
|
10796
|
-
return {
|
|
10797
|
-
id: row.id,
|
|
10798
|
-
projectId: row.project_id,
|
|
10799
|
-
executionId: row.execution_id,
|
|
10800
|
-
attempt: Number(row.attempt),
|
|
10801
|
-
callId: row.call_id,
|
|
10802
|
-
requesterGroupIds: groupRows.map((group) => group.group_id),
|
|
10803
|
-
providerId: row.provider_id,
|
|
10804
|
-
requestedModelId: row.requested_model_id,
|
|
10805
|
-
...row.response_model_id === null ? {} : { responseModelId: row.response_model_id },
|
|
10806
|
-
responseId: row.response_id,
|
|
10807
|
-
usage: {
|
|
10808
|
-
...usageFromRow2(row),
|
|
10809
|
-
...row.total_tokens === null ? {} : { totalTokens: Number(row.total_tokens) },
|
|
10810
|
-
reportingStatus: row.reporting_status
|
|
10811
|
-
},
|
|
10812
|
-
...row.raw_usage === null ? {} : { rawUsage: rawUsageFromRow(row.raw_usage) },
|
|
10813
|
-
occurredAt: new Date(row.occurred_at),
|
|
10814
|
-
recordedAt: new Date(row.recorded_at)
|
|
10815
|
-
};
|
|
10816
11952
|
}
|
|
10817
|
-
|
|
10818
|
-
const
|
|
10819
|
-
|
|
10820
|
-
|
|
10821
|
-
|
|
11953
|
+
async moveEstimateCounters(sql, reservation, directions, updatedAt) {
|
|
11954
|
+
for (const bucket of reservation.buckets) {
|
|
11955
|
+
const amount = normalizeAiLimitAmount(bucket.estimate);
|
|
11956
|
+
await this.changeCounters(sql, reservation.projectId, bucket.subject, amount, reservation.period, {
|
|
11957
|
+
...directions.reserved === undefined ? {} : { reserved: directions.reserved * amount.amount },
|
|
11958
|
+
...directions.unknown === undefined ? {} : { unknown: directions.unknown * amount.amount }
|
|
11959
|
+
}, updatedAt);
|
|
11960
|
+
}
|
|
11961
|
+
}
|
|
11962
|
+
async reconcileEstimateCounters(sql, reservation, entry, updatedAt) {
|
|
11963
|
+
for (const bucket of reservation.buckets) {
|
|
11964
|
+
const amount = normalizeAiLimitAmount(bucket.estimate);
|
|
11965
|
+
const actual = resolveAiLimitActual([entry], amount);
|
|
11966
|
+
if (actual.accountingStatus === "complete") {
|
|
11967
|
+
await this.changeCounters(sql, reservation.projectId, bucket.subject, amount, reservation.period, reservation.state === "active" ? { reserved: -amount.amount } : { unknown: -amount.amount }, updatedAt);
|
|
11968
|
+
} else if (reservation.state === "active") {
|
|
11969
|
+
await this.changeCounters(sql, reservation.projectId, bucket.subject, amount, reservation.period, { reserved: -amount.amount, unknown: amount.amount }, updatedAt);
|
|
11970
|
+
}
|
|
11971
|
+
}
|
|
11972
|
+
}
|
|
11973
|
+
async insertReservation(sql, reservation) {
|
|
11974
|
+
await sql`
|
|
11975
|
+
INSERT INTO ai_model_call_reservations (
|
|
11976
|
+
project_id, execution_id, attempt, call_id, buckets, request_key,
|
|
11977
|
+
period_start, period_end, state, usage_record_id, reserved_at, updated_at
|
|
11978
|
+
) VALUES (
|
|
11979
|
+
${reservation.projectId}, ${reservation.executionId}, ${reservation.attempt},
|
|
11980
|
+
${reservation.callId}, ${JSON.stringify(reservation.buckets)}::text::jsonb,
|
|
11981
|
+
${reservation.requestKey}, ${reservation.period.start},
|
|
11982
|
+
${reservation.period.end}, ${reservation.state}, ${reservation.usageRecordId ?? null},
|
|
11983
|
+
${reservation.reservedAt}, ${reservation.updatedAt}
|
|
10822
11984
|
)
|
|
10823
|
-
SELECT usage.* FROM requested
|
|
10824
|
-
JOIN ai_model_call_usage AS usage
|
|
10825
|
-
ON usage.project_id = ${input.projectId}
|
|
10826
|
-
AND usage.execution_id = requested.execution_id
|
|
10827
11985
|
`;
|
|
10828
11986
|
}
|
|
11987
|
+
async updateReservation(sql, reservation) {
|
|
11988
|
+
await sql`
|
|
11989
|
+
UPDATE ai_model_call_reservations
|
|
11990
|
+
SET state = ${reservation.state},
|
|
11991
|
+
usage_record_id = ${reservation.usageRecordId ?? null},
|
|
11992
|
+
updated_at = ${reservation.updatedAt}
|
|
11993
|
+
WHERE project_id = ${reservation.projectId}
|
|
11994
|
+
AND execution_id = ${reservation.executionId}
|
|
11995
|
+
AND attempt = ${reservation.attempt}
|
|
11996
|
+
AND call_id = ${reservation.callId}
|
|
11997
|
+
`;
|
|
11998
|
+
}
|
|
11999
|
+
async findReservation(sql, input) {
|
|
12000
|
+
const [row] = await sql`
|
|
12001
|
+
SELECT * FROM ai_model_call_reservations
|
|
12002
|
+
WHERE project_id = ${input.projectId}
|
|
12003
|
+
AND execution_id = ${input.executionId}
|
|
12004
|
+
AND attempt = ${input.attempt}
|
|
12005
|
+
AND call_id = ${input.callId}
|
|
12006
|
+
`;
|
|
12007
|
+
return row ? reservationFromRow(row) : null;
|
|
12008
|
+
}
|
|
12009
|
+
async accountingEntries(sql, policy, period) {
|
|
12010
|
+
const subject = aiLimitSubjectParts(policy.subject);
|
|
12011
|
+
const rows = await sql`
|
|
12012
|
+
SELECT usage.project_id, usage.id AS usage_record_id, usage.execution_id,
|
|
12013
|
+
usage.attempt, usage.call_id, usage.occurred_at,
|
|
12014
|
+
usage.total_tokens::text AS total_tokens,
|
|
12015
|
+
executions.requested_by_user_id,
|
|
12016
|
+
executions.requested_by_service_account_id,
|
|
12017
|
+
valuations.status AS valuation_status,
|
|
12018
|
+
valuations.currency AS valuation_currency,
|
|
12019
|
+
valuations.amount_nanos::text AS valuation_amount_nanos
|
|
12020
|
+
FROM ai_model_call_usage AS usage
|
|
12021
|
+
JOIN executions
|
|
12022
|
+
ON executions.project_id = usage.project_id AND executions.id = usage.execution_id
|
|
12023
|
+
LEFT JOIN ai_model_call_valuations AS valuations
|
|
12024
|
+
ON valuations.project_id = usage.project_id
|
|
12025
|
+
AND valuations.usage_record_id = usage.id
|
|
12026
|
+
WHERE usage.project_id = ${policy.projectId}
|
|
12027
|
+
AND usage.occurred_at >= ${period.start}
|
|
12028
|
+
AND usage.occurred_at < ${period.end}
|
|
12029
|
+
AND (
|
|
12030
|
+
${subject.type} = 'project'
|
|
12031
|
+
OR (
|
|
12032
|
+
${subject.type} = 'group'
|
|
12033
|
+
AND EXISTS (
|
|
12034
|
+
SELECT 1 FROM ai_model_call_usage_groups AS usage_groups
|
|
12035
|
+
WHERE usage_groups.project_id = usage.project_id
|
|
12036
|
+
AND usage_groups.usage_record_id = usage.id
|
|
12037
|
+
AND usage_groups.group_id = ${subject.id}
|
|
12038
|
+
)
|
|
12039
|
+
)
|
|
12040
|
+
OR (${subject.type} = 'user' AND executions.requested_by_user_id = ${subject.id})
|
|
12041
|
+
OR (
|
|
12042
|
+
${subject.type} = 'serviceAccount'
|
|
12043
|
+
AND executions.requested_by_service_account_id = ${subject.id}
|
|
12044
|
+
)
|
|
12045
|
+
)
|
|
12046
|
+
ORDER BY usage.occurred_at, usage.id
|
|
12047
|
+
`;
|
|
12048
|
+
return rows.map((row) => accountingEntryFromRow(row, policy.subject));
|
|
12049
|
+
}
|
|
12050
|
+
async requireAccountingEntry(sql, projectId, usageRecordId) {
|
|
12051
|
+
const [row] = await sql`
|
|
12052
|
+
SELECT usage.project_id, usage.id AS usage_record_id, usage.execution_id,
|
|
12053
|
+
usage.attempt, usage.call_id, usage.occurred_at,
|
|
12054
|
+
usage.total_tokens::text AS total_tokens,
|
|
12055
|
+
executions.requested_by_user_id,
|
|
12056
|
+
executions.requested_by_service_account_id,
|
|
12057
|
+
valuations.status AS valuation_status,
|
|
12058
|
+
valuations.currency AS valuation_currency,
|
|
12059
|
+
valuations.amount_nanos::text AS valuation_amount_nanos
|
|
12060
|
+
FROM ai_model_call_usage AS usage
|
|
12061
|
+
JOIN executions
|
|
12062
|
+
ON executions.project_id = usage.project_id AND executions.id = usage.execution_id
|
|
12063
|
+
LEFT JOIN ai_model_call_valuations AS valuations
|
|
12064
|
+
ON valuations.project_id = usage.project_id
|
|
12065
|
+
AND valuations.usage_record_id = usage.id
|
|
12066
|
+
WHERE usage.project_id = ${projectId} AND usage.id = ${usageRecordId}
|
|
12067
|
+
`;
|
|
12068
|
+
if (row) {
|
|
12069
|
+
const groups = await sql`
|
|
12070
|
+
SELECT group_id FROM ai_model_call_usage_groups
|
|
12071
|
+
WHERE project_id = ${projectId} AND usage_record_id = ${usageRecordId}
|
|
12072
|
+
ORDER BY group_id
|
|
12073
|
+
`;
|
|
12074
|
+
return accountingEntryFromRow(row, undefined, groups.map((group) => group.group_id));
|
|
12075
|
+
}
|
|
12076
|
+
throw new AiLimitStorageError("missing_usage_record", `[SixbPg] AI usage record '${usageRecordId}' does not exist in project '${projectId}'.`);
|
|
12077
|
+
}
|
|
10829
12078
|
async assertExecutionExists(sql, projectId, executionId) {
|
|
10830
|
-
const [
|
|
12079
|
+
const [row] = await sql`
|
|
10831
12080
|
SELECT EXISTS (
|
|
10832
12081
|
SELECT 1 FROM executions WHERE project_id = ${projectId} AND id = ${executionId}
|
|
10833
12082
|
) AS exists
|
|
10834
12083
|
`;
|
|
10835
|
-
if (
|
|
12084
|
+
if (row?.exists)
|
|
10836
12085
|
return;
|
|
10837
|
-
throw new
|
|
12086
|
+
throw new AiLimitStorageError("missing_execution", `[SixbPg] AI limit execution '${executionId}' does not exist in project '${projectId}'.`);
|
|
10838
12087
|
}
|
|
10839
12088
|
}
|
|
10840
|
-
function
|
|
10841
|
-
|
|
10842
|
-
|
|
10843
|
-
|
|
10844
|
-
|
|
10845
|
-
|
|
10846
|
-
|
|
10847
|
-
|
|
10848
|
-
|
|
10849
|
-
|
|
10850
|
-
|
|
10851
|
-
|
|
10852
|
-
|
|
10853
|
-
|
|
12089
|
+
function policyFromRow(row) {
|
|
12090
|
+
return {
|
|
12091
|
+
id: row.id,
|
|
12092
|
+
projectId: row.project_id,
|
|
12093
|
+
subject: aiLimitSubjectFromParts(row.subject_type, row.subject_id),
|
|
12094
|
+
limit: aiLimitQuantityFromAmount({
|
|
12095
|
+
meter: row.meter,
|
|
12096
|
+
currency: row.currency,
|
|
12097
|
+
amount: BigInt(row.limit_amount)
|
|
12098
|
+
}),
|
|
12099
|
+
period: row.period_kind,
|
|
12100
|
+
enabled: row.enabled,
|
|
12101
|
+
createdAt: new Date(row.created_at),
|
|
12102
|
+
updatedAt: new Date(row.updated_at)
|
|
12103
|
+
};
|
|
10854
12104
|
}
|
|
10855
|
-
function
|
|
12105
|
+
function policyStatus(policy, period, row, existingGroups) {
|
|
12106
|
+
const limit = normalizeAiLimitAmount(policy.limit);
|
|
12107
|
+
const actual = BigInt(row.actual_amount);
|
|
12108
|
+
const reserved = BigInt(row.reserved_amount);
|
|
12109
|
+
const unknown = BigInt(row.unknown_amount);
|
|
12110
|
+
const consumed = actual + reserved + unknown;
|
|
12111
|
+
const quantity = (amount) => aiLimitQuantityFromAmount({ ...limit, amount });
|
|
10856
12112
|
return {
|
|
10857
|
-
|
|
10858
|
-
|
|
10859
|
-
|
|
10860
|
-
|
|
10861
|
-
|
|
10862
|
-
|
|
10863
|
-
|
|
12113
|
+
policy,
|
|
12114
|
+
period: structuredClone(period),
|
|
12115
|
+
consumption: {
|
|
12116
|
+
actual: quantity(actual),
|
|
12117
|
+
reserved: quantity(reserved),
|
|
12118
|
+
unknown: quantity(unknown),
|
|
12119
|
+
remaining: quantity(limit.amount > consumed ? limit.amount - consumed : 0n)
|
|
12120
|
+
},
|
|
12121
|
+
accountingStatus: row.accounting_status,
|
|
12122
|
+
exhausted: consumed >= limit.amount,
|
|
12123
|
+
orphaned: policy.subject.type === "group" && existingGroups !== undefined && !existingGroups.has(policy.subject.id)
|
|
10864
12124
|
};
|
|
10865
12125
|
}
|
|
10866
|
-
function
|
|
12126
|
+
function accountingEntryFromRow(row, subject, requesterGroupIds = subject?.type === "group" ? [subject.id] : []) {
|
|
12127
|
+
const totalTokens = row.total_tokens === null ? undefined : Number(row.total_tokens);
|
|
12128
|
+
const requester = row.requested_by_user_id ? { type: "user", id: row.requested_by_user_id } : row.requested_by_service_account_id ? { type: "serviceAccount", id: row.requested_by_service_account_id } : undefined;
|
|
12129
|
+
return {
|
|
12130
|
+
projectId: row.project_id,
|
|
12131
|
+
usageRecordId: row.usage_record_id,
|
|
12132
|
+
executionId: row.execution_id,
|
|
12133
|
+
attempt: Number(row.attempt),
|
|
12134
|
+
callId: row.call_id,
|
|
12135
|
+
occurredAt: new Date(row.occurred_at),
|
|
12136
|
+
...totalTokens !== undefined && Number.isSafeInteger(totalTokens) ? { totalTokens } : {},
|
|
12137
|
+
requesterGroupIds,
|
|
12138
|
+
...requester ? { requester } : {},
|
|
12139
|
+
...row.valuation_status === "rated" && row.valuation_currency !== null && row.valuation_amount_nanos !== null ? {
|
|
12140
|
+
valuation: {
|
|
12141
|
+
status: "rated",
|
|
12142
|
+
money: {
|
|
12143
|
+
currency: row.valuation_currency,
|
|
12144
|
+
amountNanos: row.valuation_amount_nanos
|
|
12145
|
+
}
|
|
12146
|
+
}
|
|
12147
|
+
} : row.valuation_status === "unpriceable" ? { valuation: { status: "unavailable" } } : {}
|
|
12148
|
+
};
|
|
12149
|
+
}
|
|
12150
|
+
function reservationFromRow(row) {
|
|
12151
|
+
const buckets = jsonValue(row.buckets);
|
|
12152
|
+
return {
|
|
12153
|
+
projectId: row.project_id,
|
|
12154
|
+
executionId: row.execution_id,
|
|
12155
|
+
attempt: Number(row.attempt),
|
|
12156
|
+
callId: row.call_id,
|
|
12157
|
+
buckets: buckets.map((bucket) => ({
|
|
12158
|
+
subject: structuredClone(bucket.subject),
|
|
12159
|
+
estimate: normalizeAiLimitQuantities([bucket.estimate], "stored bucket")[0]
|
|
12160
|
+
})),
|
|
12161
|
+
requestKey: row.request_key,
|
|
12162
|
+
period: {
|
|
12163
|
+
kind: "calendarMonth",
|
|
12164
|
+
start: new Date(row.period_start),
|
|
12165
|
+
end: new Date(row.period_end),
|
|
12166
|
+
resetAt: new Date(row.period_end)
|
|
12167
|
+
},
|
|
12168
|
+
state: row.state,
|
|
12169
|
+
...row.usage_record_id === null ? {} : { usageRecordId: row.usage_record_id },
|
|
12170
|
+
reservedAt: new Date(row.reserved_at),
|
|
12171
|
+
updatedAt: new Date(row.updated_at)
|
|
12172
|
+
};
|
|
12173
|
+
}
|
|
12174
|
+
function jsonValue(value) {
|
|
10867
12175
|
return typeof value === "string" ? JSON.parse(value) : value;
|
|
10868
12176
|
}
|
|
10869
|
-
function
|
|
10870
|
-
if (
|
|
10871
|
-
throw new TypeError("[Sixb] AI
|
|
12177
|
+
function assertSameDimension(previous, next) {
|
|
12178
|
+
if (aiLimitAmountKey(normalizeAiLimitAmount(previous)) !== aiLimitAmountKey(normalizeAiLimitAmount(next))) {
|
|
12179
|
+
throw new TypeError("[Sixb] AI limit policy meter and cost currency are immutable.");
|
|
10872
12180
|
}
|
|
10873
12181
|
}
|
|
12182
|
+
function validatePolicyIdentity(input) {
|
|
12183
|
+
assertNonBlank2(input.projectId, "projectId");
|
|
12184
|
+
assertNonBlank2(input.id, "policy id");
|
|
12185
|
+
}
|
|
12186
|
+
function duplicatePolicy(policy) {
|
|
12187
|
+
return new AiLimitStorageError("duplicate_policy", `[SixbPg] An AI limit policy already exists for policy ID or subject-meter dimension '${policy.id}'.`);
|
|
12188
|
+
}
|
|
12189
|
+
function missingReservation(input) {
|
|
12190
|
+
return new AiLimitStorageError("missing_reservation", `[SixbPg] AI model-call reservation '${input.callId}' does not exist for execution '${input.executionId}' attempt ${input.attempt}.`);
|
|
12191
|
+
}
|
|
12192
|
+
function reconciliationConflict(callId) {
|
|
12193
|
+
return new AiLimitStorageError("reconciliation_conflict", `[SixbPg] AI model-call reservation '${callId}' was reconciled with a different usage record.`);
|
|
12194
|
+
}
|
|
12195
|
+
function invalidState(reservation, operation) {
|
|
12196
|
+
return new AiLimitStorageError("invalid_reservation_state", `[SixbPg] Cannot ${operation} AI model-call reservation '${reservation.callId}' in state '${reservation.state}'.`);
|
|
12197
|
+
}
|
|
10874
12198
|
|
|
10875
12199
|
// src/pg-client.ts
|
|
10876
12200
|
import postgres from "postgres";
|
|
@@ -14724,15 +16048,15 @@ class PgWorkflowAgentNodeRunStorage {
|
|
|
14724
16048
|
projectId: input.projectId,
|
|
14725
16049
|
executionId: input.executionId,
|
|
14726
16050
|
nodeRunId: input.nodeRunId,
|
|
14727
|
-
|
|
16051
|
+
actorId: input.actorId,
|
|
14728
16052
|
workflowExecutionId: parent.execution_id
|
|
14729
16053
|
});
|
|
14730
16054
|
try {
|
|
14731
16055
|
const [row] = await this.sql`
|
|
14732
16056
|
INSERT INTO workflow_agent_node_runs (
|
|
14733
|
-
project_id, node_run_id, execution_id,
|
|
16057
|
+
project_id, node_run_id, execution_id, actor_id, status, prompt, attempt, created_at
|
|
14734
16058
|
) VALUES (
|
|
14735
|
-
${input.projectId}, ${input.nodeRunId}, ${input.executionId}, ${input.
|
|
16059
|
+
${input.projectId}, ${input.nodeRunId}, ${input.executionId}, ${input.actorId},
|
|
14736
16060
|
${"queued"}, ${input.prompt}, ${0}, ${input.createdAt ?? new Date}
|
|
14737
16061
|
) RETURNING *
|
|
14738
16062
|
`;
|
|
@@ -14849,9 +16173,9 @@ class PgWorkflowAgentNodeRunStorage {
|
|
|
14849
16173
|
const whereClauses = ["project_id = $1"];
|
|
14850
16174
|
const params = [input.projectId];
|
|
14851
16175
|
let index = 2;
|
|
14852
|
-
if (input.
|
|
14853
|
-
whereClauses.push(`
|
|
14854
|
-
params.push(input.
|
|
16176
|
+
if (input.actorId) {
|
|
16177
|
+
whereClauses.push(`actor_id = $${index++}`);
|
|
16178
|
+
params.push(input.actorId);
|
|
14855
16179
|
}
|
|
14856
16180
|
index = appendRunListFilters(whereClauses, params, index, input);
|
|
14857
16181
|
const result = await queryRunList({
|
|
@@ -15128,7 +16452,7 @@ function rowToWorkflowAgentNodeRunRecord(row) {
|
|
|
15128
16452
|
projectId: row.project_id,
|
|
15129
16453
|
nodeRunId: row.node_run_id,
|
|
15130
16454
|
executionId: row.execution_id,
|
|
15131
|
-
|
|
16455
|
+
actorId: row.actor_id,
|
|
15132
16456
|
status: row.status,
|
|
15133
16457
|
prompt: row.prompt,
|
|
15134
16458
|
...row.model_id ? { modelId: row.model_id } : {},
|
|
@@ -15193,6 +16517,7 @@ class PostgresStorage {
|
|
|
15193
16517
|
agents;
|
|
15194
16518
|
aiUsage;
|
|
15195
16519
|
aiCosts;
|
|
16520
|
+
aiLimits;
|
|
15196
16521
|
actionRuns;
|
|
15197
16522
|
pipelineRuns;
|
|
15198
16523
|
workflowRuns;
|
|
@@ -15243,6 +16568,7 @@ class PostgresStorage {
|
|
|
15243
16568
|
this.agents = createAgentOperationScope(stores.agents, scope);
|
|
15244
16569
|
this.aiUsage = createOperationScopedFacade(stores.aiUsage, scope);
|
|
15245
16570
|
this.aiCosts = createOperationScopedFacade(stores.aiCosts, scope);
|
|
16571
|
+
this.aiLimits = createOperationScopedFacade(stores.aiLimits, scope);
|
|
15246
16572
|
this.actionRuns = createOperationScopedFacade(stores.actionRuns, scope);
|
|
15247
16573
|
this.pipelineRuns = createOperationScopedFacade(stores.pipelineRuns, scope);
|
|
15248
16574
|
this.workflowRuns = createWorkflowRunOperationScope(stores.workflowRuns, scope);
|
|
@@ -15278,6 +16604,7 @@ class PostgresStorage {
|
|
|
15278
16604
|
transactionContext.materializations.assertCommittable();
|
|
15279
16605
|
return result;
|
|
15280
16606
|
} finally {
|
|
16607
|
+
await txStorage.aiLimits.settleOperations();
|
|
15281
16608
|
transactionContext.active = false;
|
|
15282
16609
|
txStorage.ontology.deactivateSessions();
|
|
15283
16610
|
transactionContext.materializations.deactivate();
|
|
@@ -15332,6 +16659,7 @@ function createPostgresStores(sql, options) {
|
|
|
15332
16659
|
agents: new PgAgentStorage({ sql, executions }),
|
|
15333
16660
|
aiUsage: new PgAiUsageStorage(sql),
|
|
15334
16661
|
aiCosts: new PgAiCostStorage(sql),
|
|
16662
|
+
aiLimits: new PgAiLimitStorage(sql),
|
|
15335
16663
|
actionRuns: new PgActionRunStorage(sql, executions),
|
|
15336
16664
|
pipelineRuns: new PgPipelineRunStorage(sql, executions),
|
|
15337
16665
|
workflowRuns: new PgWorkflowRunStorage(sql, executions),
|
|
@@ -15358,4 +16686,4 @@ export {
|
|
|
15358
16686
|
PostgresStorage
|
|
15359
16687
|
};
|
|
15360
16688
|
|
|
15361
|
-
//# debugId=
|
|
16689
|
+
//# debugId=A5AA822A44042C7A64756E2164756E21
|