@sixb/pg 0.1.5 → 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 +1804 -464
- package/dist/index.js.map +19 -16
- package/dist/migrations.d.ts.map +1 -1
- package/dist/ontology-storage/materialization-state.d.ts +7 -3
- package/dist/ontology-storage/materialization-state.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/ontology-storage/materialization-state.ts +45 -29
- package/src/ontology-storage/materializations.ts +3 -3
- 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
|
|
|
@@ -6936,7 +7518,12 @@ class PgMaterializationStateReader {
|
|
|
6936
7518
|
return stored ? [stored] : [];
|
|
6937
7519
|
});
|
|
6938
7520
|
}
|
|
6939
|
-
async
|
|
7521
|
+
async linkScopeRevisions(scopes) {
|
|
7522
|
+
if (scopes.length === 0)
|
|
7523
|
+
return [];
|
|
7524
|
+
return (await this.readEffectiveLinkScopes(scopes)).revisions;
|
|
7525
|
+
}
|
|
7526
|
+
async linkSlotStates(scopes) {
|
|
6940
7527
|
if (scopes.length === 0)
|
|
6941
7528
|
return [];
|
|
6942
7529
|
const requested = scopes.map(({ source, linkId }) => ({
|
|
@@ -6945,10 +7532,6 @@ class PgMaterializationStateReader {
|
|
|
6945
7532
|
source_id: source.primaryId,
|
|
6946
7533
|
link_id: linkId
|
|
6947
7534
|
}));
|
|
6948
|
-
const accumulators = new Map(scopes.map(({ source, linkId }) => {
|
|
6949
|
-
const key = linkScopeSortKey(source, linkId);
|
|
6950
|
-
return [key, startScopeAccumulator(source, linkId)];
|
|
6951
|
-
}));
|
|
6952
7535
|
const requestedParameter = jsonParameter2(this.sql, requested);
|
|
6953
7536
|
const [sourceRows, slotOverrideRows] = await Promise.all([
|
|
6954
7537
|
this.sql`
|
|
@@ -6994,17 +7577,40 @@ class PgMaterializationStateReader {
|
|
|
6994
7577
|
const stored = storedLinkSlotOverride(row);
|
|
6995
7578
|
return [linkScopeSortKey(stored.ref.source, stored.ref.linkId), stored];
|
|
6996
7579
|
}));
|
|
6997
|
-
const effective =
|
|
6998
|
-
|
|
6999
|
-
|
|
7000
|
-
|
|
7001
|
-
|
|
7002
|
-
|
|
7003
|
-
|
|
7004
|
-
|
|
7005
|
-
|
|
7006
|
-
|
|
7007
|
-
|
|
7580
|
+
const { revisions, effective } = await this.readEffectiveLinkScopes(scopes);
|
|
7581
|
+
return scopes.map(({ source, linkId }, index) => {
|
|
7582
|
+
const key = linkScopeSortKey(source, linkId);
|
|
7583
|
+
return {
|
|
7584
|
+
...revisions[index],
|
|
7585
|
+
sourceAssertion: sources.get(key) ?? null,
|
|
7586
|
+
override: overrides.get(key) ?? null,
|
|
7587
|
+
effective: effective.get(key) ?? null
|
|
7588
|
+
};
|
|
7589
|
+
});
|
|
7590
|
+
}
|
|
7591
|
+
async readEffectiveLinkScopes(scopes) {
|
|
7592
|
+
const requested = scopes.map(({ source, linkId }) => ({
|
|
7593
|
+
scope_sort_key: linkScopeSortKey(source, linkId),
|
|
7594
|
+
source_type_id: source.objectTypeId,
|
|
7595
|
+
source_id: source.primaryId,
|
|
7596
|
+
link_id: linkId
|
|
7597
|
+
}));
|
|
7598
|
+
const requestedParameter = jsonParameter2(this.sql, requested);
|
|
7599
|
+
const accumulators = new Map(scopes.map(({ source, linkId }) => {
|
|
7600
|
+
const key = linkScopeSortKey(source, linkId);
|
|
7601
|
+
return [key, startScopeAccumulator(source, linkId)];
|
|
7602
|
+
}));
|
|
7603
|
+
const effective = new Map;
|
|
7604
|
+
let scopeCursor = null;
|
|
7605
|
+
let linkCursor = null;
|
|
7606
|
+
while (true) {
|
|
7607
|
+
const rows = await this.sql`
|
|
7608
|
+
WITH requested AS (
|
|
7609
|
+
SELECT DISTINCT * FROM jsonb_to_recordset(${requestedParameter})
|
|
7610
|
+
AS requested_values(
|
|
7611
|
+
scope_sort_key TEXT, source_type_id TEXT, source_id TEXT, link_id TEXT
|
|
7612
|
+
)
|
|
7613
|
+
), selected AS (
|
|
7008
7614
|
SELECT requested.scope_sort_key, links.*,
|
|
7009
7615
|
${this.sql.unsafe(linkSortExpression("links"))} AS link_sort_key
|
|
7010
7616
|
FROM links
|
|
@@ -7032,22 +7638,10 @@ class PgMaterializationStateReader {
|
|
|
7032
7638
|
scopeCursor = last.scope_sort_key;
|
|
7033
7639
|
linkCursor = last.link_sort_key;
|
|
7034
7640
|
}
|
|
7035
|
-
return
|
|
7036
|
-
|
|
7037
|
-
|
|
7038
|
-
|
|
7039
|
-
sourceAssertion: sources.get(key) ?? null,
|
|
7040
|
-
override: overrides.get(key) ?? null,
|
|
7041
|
-
effective: effective.get(key) ?? null
|
|
7042
|
-
};
|
|
7043
|
-
});
|
|
7044
|
-
}
|
|
7045
|
-
async linkScope(source, linkId) {
|
|
7046
|
-
const [scope] = await this.linkScopes([{ source, linkId }]);
|
|
7047
|
-
if (!scope) {
|
|
7048
|
-
throw new MaterializationConflictError2("effective-state", "Link scope lookup returned no row.");
|
|
7049
|
-
}
|
|
7050
|
-
return scope;
|
|
7641
|
+
return {
|
|
7642
|
+
revisions: scopes.map(({ source, linkId }) => finishScopeAccumulator(accumulators.get(linkScopeSortKey(source, linkId)))),
|
|
7643
|
+
effective
|
|
7644
|
+
};
|
|
7051
7645
|
}
|
|
7052
7646
|
async* incidentLinks(objects, pageRows) {
|
|
7053
7647
|
if (objects.length === 0)
|
|
@@ -8453,9 +9047,9 @@ class PgOntologyMaterializationStorage {
|
|
|
8453
9047
|
for (const expected of input.expected.links) {
|
|
8454
9048
|
this.assertLink(linkRevisions.get(linkRefKey3(expected.ref)), expected);
|
|
8455
9049
|
}
|
|
8456
|
-
const
|
|
9050
|
+
const linkScopeRevisions = await reader.linkScopeRevisions(input.expected.linkScopes);
|
|
8457
9051
|
for (const [index, expected] of input.expected.linkScopes.entries()) {
|
|
8458
|
-
if (
|
|
9052
|
+
if (linkScopeRevisions[index]?.fingerprint !== expected.fingerprint) {
|
|
8459
9053
|
throw new MaterializationConflictError5("effective-state", `Expected link scope changed for ${expected.source.objectTypeId}:${expected.source.primaryId}.${expected.linkId}.`);
|
|
8460
9054
|
}
|
|
8461
9055
|
}
|
|
@@ -8502,7 +9096,7 @@ class PgOntologyMaterializationStorage {
|
|
|
8502
9096
|
yield {
|
|
8503
9097
|
objects: [],
|
|
8504
9098
|
links: [],
|
|
8505
|
-
linkScopes: await reader.
|
|
9099
|
+
linkScopes: await reader.linkSlotStates(scopes.slice(offset, offset + input.pageRows)),
|
|
8506
9100
|
points: []
|
|
8507
9101
|
};
|
|
8508
9102
|
}
|
|
@@ -10153,90 +10747,429 @@ import {
|
|
|
10153
10747
|
aiModelCallCostMatchesUsage,
|
|
10154
10748
|
buildAiAccountingOverviewFromFragments,
|
|
10155
10749
|
normalizeAiAccountingQuery,
|
|
10156
|
-
normalizeAiModelCallAccountingQuery,
|
|
10750
|
+
normalizeAiModelCallAccountingQuery as normalizeAiModelCallAccountingQuery2,
|
|
10157
10751
|
normalizeAiModelCallCostRecord,
|
|
10158
10752
|
parseAiModelCallCostDetails,
|
|
10159
10753
|
toAccountingItem
|
|
10160
10754
|
} from "@sixb/core/internal/ai-cost-storage-provider";
|
|
10755
|
+
import { normalizeModelProviderIds as normalizeModelProviderIds2 } from "@sixb/core/models";
|
|
10161
10756
|
import {
|
|
10162
10757
|
AiCostStorageError
|
|
10163
10758
|
} from "@sixb/core/storage";
|
|
10164
|
-
|
|
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 {
|
|
10165
10864
|
sql;
|
|
10166
10865
|
constructor(sql) {
|
|
10167
10866
|
this.sql = sql;
|
|
10168
10867
|
}
|
|
10169
|
-
async
|
|
10170
|
-
const record =
|
|
10171
|
-
|
|
10172
|
-
const
|
|
10173
|
-
if (
|
|
10174
|
-
|
|
10175
|
-
|
|
10176
|
-
if (await findCost(tx, record.projectId, record.usageRecordId))
|
|
10177
|
-
return;
|
|
10178
|
-
const rated = record.status === "rated" ? record : undefined;
|
|
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);
|
|
10179
10875
|
const inserted = await tx`
|
|
10180
|
-
INSERT INTO
|
|
10181
|
-
project_id,
|
|
10182
|
-
|
|
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
|
|
10183
10900
|
) VALUES (
|
|
10184
|
-
${record.projectId},
|
|
10185
|
-
${record.
|
|
10186
|
-
${record.
|
|
10187
|
-
${
|
|
10188
|
-
${record.
|
|
10189
|
-
${
|
|
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}
|
|
10190
10924
|
)
|
|
10191
10925
|
ON CONFLICT DO NOTHING
|
|
10192
|
-
RETURNING
|
|
10926
|
+
RETURNING id
|
|
10193
10927
|
`;
|
|
10194
|
-
if (inserted.length === 0
|
|
10195
|
-
|
|
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
|
+
`;
|
|
10196
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]
|
|
10197
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;
|
|
10198
10963
|
}
|
|
10199
10964
|
async summarizeExecutions(input) {
|
|
10200
|
-
|
|
10201
|
-
for (const
|
|
10202
|
-
|
|
10965
|
+
assertNonBlankProjectId(input.projectId);
|
|
10966
|
+
for (const executionId of input.executionIds)
|
|
10967
|
+
assertAiUsageExecutionId(executionId);
|
|
10203
10968
|
if (input.executionIds.length === 0)
|
|
10204
10969
|
return [];
|
|
10205
|
-
const
|
|
10206
|
-
|
|
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`
|
|
10207
11026
|
WITH requested(execution_id) AS (
|
|
10208
|
-
SELECT
|
|
11027
|
+
SELECT * FROM unnest(${sql.array(executionIds)}::text[])
|
|
10209
11028
|
)
|
|
10210
|
-
SELECT usage
|
|
10211
|
-
cost.status,
|
|
10212
|
-
cost.currency,
|
|
10213
|
-
cost.amount_nanos
|
|
10214
|
-
FROM requested
|
|
11029
|
+
SELECT usage.* FROM requested
|
|
10215
11030
|
JOIN ai_model_call_usage AS usage
|
|
10216
|
-
ON usage.project_id = ${input.projectId}
|
|
10217
|
-
|
|
10218
|
-
ON cost.project_id = usage.project_id AND cost.usage_record_id = usage.id
|
|
11031
|
+
ON usage.project_id = ${input.projectId}
|
|
11032
|
+
AND usage.execution_id = requested.execution_id
|
|
10219
11033
|
`;
|
|
10220
|
-
return aggregateSummaryRows(input.executionIds, rows);
|
|
10221
11034
|
}
|
|
10222
|
-
async
|
|
10223
|
-
const
|
|
10224
|
-
|
|
10225
|
-
|
|
10226
|
-
|
|
10227
|
-
|
|
10228
|
-
|
|
10229
|
-
|
|
10230
|
-
|
|
10231
|
-
|
|
10232
|
-
|
|
10233
|
-
|
|
10234
|
-
|
|
10235
|
-
|
|
10236
|
-
|
|
10237
|
-
|
|
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,
|
|
11110
|
+
currency, amount_nanos, reason, details, rated_at
|
|
11111
|
+
) VALUES (
|
|
11112
|
+
${record.projectId}, ${record.usageRecordId}, ${record.status},
|
|
11113
|
+
${record.billingIdentity?.providerId ?? null},
|
|
11114
|
+
${record.billingIdentity?.modelId ?? null}, ${rated?.money.currency ?? null},
|
|
11115
|
+
${rated?.money.amountNanos ?? null},
|
|
11116
|
+
${record.status === "unpriceable" ? record.reason : null},
|
|
11117
|
+
${JSON.stringify(aiModelCallCostDetails(record))}::text::jsonb, ${record.ratedAt}
|
|
11118
|
+
)
|
|
11119
|
+
ON CONFLICT DO NOTHING
|
|
11120
|
+
RETURNING usage_record_id
|
|
11121
|
+
`;
|
|
11122
|
+
if (inserted.length === 0 && !await findCost(tx, record.projectId, record.usageRecordId)) {
|
|
11123
|
+
throw new Error("[SixbPg] AI valuation insert conflicted without a visible row.");
|
|
11124
|
+
}
|
|
11125
|
+
});
|
|
11126
|
+
}
|
|
11127
|
+
async summarizeExecutions(input) {
|
|
11128
|
+
assertNonBlank(input.projectId, "cost summary projectId");
|
|
11129
|
+
for (const id of input.executionIds)
|
|
11130
|
+
assertNonBlank(id, "cost summary executionId");
|
|
11131
|
+
if (input.executionIds.length === 0)
|
|
11132
|
+
return [];
|
|
11133
|
+
const ids = [...new Set(input.executionIds)];
|
|
11134
|
+
const rows = await this.sql`
|
|
11135
|
+
WITH requested(execution_id) AS (
|
|
11136
|
+
SELECT DISTINCT * FROM unnest(${this.sql.array(ids)}::text[])
|
|
11137
|
+
)
|
|
11138
|
+
SELECT usage.execution_id,
|
|
11139
|
+
cost.status,
|
|
11140
|
+
cost.currency,
|
|
11141
|
+
cost.amount_nanos
|
|
11142
|
+
FROM requested
|
|
11143
|
+
JOIN ai_model_call_usage AS usage
|
|
11144
|
+
ON usage.project_id = ${input.projectId} AND usage.execution_id = requested.execution_id
|
|
11145
|
+
LEFT JOIN ai_model_call_valuations AS cost
|
|
11146
|
+
ON cost.project_id = usage.project_id AND cost.usage_record_id = usage.id
|
|
11147
|
+
`;
|
|
11148
|
+
return aggregateSummaryRows(input.executionIds, rows);
|
|
11149
|
+
}
|
|
11150
|
+
async queryProjectOverview(input) {
|
|
11151
|
+
const query = normalizeAiAccountingQuery(input);
|
|
11152
|
+
const providerId = query.providerId ?? null;
|
|
11153
|
+
const modelId = query.modelId ?? null;
|
|
11154
|
+
const rows = await this.sql`
|
|
11155
|
+
WITH filtered AS (
|
|
11156
|
+
SELECT usage.provider_id,
|
|
11157
|
+
usage.requested_model_id,
|
|
11158
|
+
usage.occurred_at,
|
|
11159
|
+
usage.input_tokens,
|
|
11160
|
+
usage.output_tokens,
|
|
11161
|
+
usage.uncached_input_tokens,
|
|
11162
|
+
usage.cache_read_input_tokens,
|
|
11163
|
+
usage.cache_write_input_tokens,
|
|
11164
|
+
usage.text_output_tokens,
|
|
11165
|
+
usage.reasoning_output_tokens,
|
|
10238
11166
|
usage.reporting_status,
|
|
10239
|
-
|
|
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,
|
|
10240
11173
|
workflow_node.workflow_id AS attribution_workflow_id,
|
|
10241
11174
|
cost.status AS cost_status,
|
|
10242
11175
|
cost.currency AS cost_currency,
|
|
@@ -10271,11 +11204,12 @@ class PgAiCostStorage {
|
|
|
10271
11204
|
SELECT 'totals'::text AS dimension,
|
|
10272
11205
|
NULL::text AS key_1,
|
|
10273
11206
|
NULL::text AS key_2,
|
|
11207
|
+
NULL::text AS key_3,
|
|
10274
11208
|
NULL::timestamptz AS bucket_start,
|
|
10275
11209
|
classified.*
|
|
10276
11210
|
FROM classified
|
|
10277
11211
|
UNION ALL
|
|
10278
|
-
SELECT 'series', NULL, NULL,
|
|
11212
|
+
SELECT 'series', NULL, NULL, NULL,
|
|
10279
11213
|
CASE ${input.bucket}::text
|
|
10280
11214
|
WHEN 'hour' THEN
|
|
10281
11215
|
date_trunc('hour', occurred_at AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'
|
|
@@ -10286,17 +11220,18 @@ class PgAiCostStorage {
|
|
|
10286
11220
|
classified.*
|
|
10287
11221
|
FROM classified
|
|
10288
11222
|
UNION ALL
|
|
10289
|
-
SELECT 'model', provider_id, requested_model_id, NULL, classified.* FROM classified
|
|
11223
|
+
SELECT 'model', provider_id, requested_model_id, NULL, NULL, classified.* FROM classified
|
|
10290
11224
|
UNION ALL
|
|
10291
|
-
SELECT 'agent',
|
|
10292
|
-
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
|
|
10293
11227
|
UNION ALL
|
|
10294
|
-
SELECT 'workflow', attribution_workflow_id, NULL, NULL, classified.*
|
|
11228
|
+
SELECT 'workflow', attribution_workflow_id, NULL, NULL, NULL, classified.*
|
|
10295
11229
|
FROM classified WHERE attribution_workflow_id IS NOT NULL
|
|
10296
11230
|
)
|
|
10297
11231
|
SELECT dimension,
|
|
10298
11232
|
key_1,
|
|
10299
11233
|
key_2,
|
|
11234
|
+
key_3,
|
|
10300
11235
|
bucket_start,
|
|
10301
11236
|
amount_currency,
|
|
10302
11237
|
COUNT(*)::text AS model_call_count,
|
|
@@ -10323,12 +11258,12 @@ class PgAiCostStorage {
|
|
|
10323
11258
|
AS unvalued_call_count,
|
|
10324
11259
|
SUM(amount_nanos)::text AS amount_nanos
|
|
10325
11260
|
FROM expanded
|
|
10326
|
-
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
|
|
10327
11262
|
`;
|
|
10328
11263
|
return buildAiAccountingOverviewFromFragments(input, rows.map(aggregateFragmentFromRow));
|
|
10329
11264
|
}
|
|
10330
11265
|
async listModelCalls(input) {
|
|
10331
|
-
const query =
|
|
11266
|
+
const query = normalizeAiModelCallAccountingQuery2(input);
|
|
10332
11267
|
const providerId = query.providerId ?? null;
|
|
10333
11268
|
const modelId = query.modelId ?? null;
|
|
10334
11269
|
const executionId = query.executionId ?? null;
|
|
@@ -10354,13 +11289,15 @@ class PgAiCostStorage {
|
|
|
10354
11289
|
const rows = await this.sql`
|
|
10355
11290
|
SELECT usage.*,
|
|
10356
11291
|
CASE
|
|
10357
|
-
WHEN direct_agent.
|
|
11292
|
+
WHEN direct_agent.kind = 'conversation' THEN 'agent'
|
|
11293
|
+
WHEN direct_agent.kind = 'subagent' THEN 'subagent'
|
|
10358
11294
|
WHEN workflow_agent.node_run_id IS NOT NULL THEN 'workflowAgent'
|
|
10359
11295
|
ELSE NULL
|
|
10360
11296
|
END AS attribution_kind,
|
|
10361
|
-
|
|
11297
|
+
workflow_node.node_id AS attribution_agent_step_id,
|
|
10362
11298
|
direct_agent.id AS attribution_agent_run_id,
|
|
10363
11299
|
direct_agent.thread_id AS attribution_thread_id,
|
|
11300
|
+
direct_agent.parent_run_id AS attribution_parent_run_id,
|
|
10364
11301
|
workflow_agent.node_run_id AS attribution_node_run_id,
|
|
10365
11302
|
workflow_node.workflow_id AS attribution_workflow_id,
|
|
10366
11303
|
workflow_node.workflow_run_id AS attribution_workflow_run_id,
|
|
@@ -10426,10 +11363,13 @@ async function findUsage(sql, projectId, id) {
|
|
|
10426
11363
|
const [row] = await sql`
|
|
10427
11364
|
SELECT * FROM ai_model_call_usage WHERE project_id = ${projectId} AND id = ${id}
|
|
10428
11365
|
`;
|
|
10429
|
-
return row ?
|
|
11366
|
+
return row ? usageFromRow2(row, []) : null;
|
|
10430
11367
|
}
|
|
10431
11368
|
function aggregateFragmentFromRow(row) {
|
|
10432
|
-
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)) } : {};
|
|
10433
11373
|
return {
|
|
10434
11374
|
dimension: row.dimension,
|
|
10435
11375
|
...dimensionKeys,
|
|
@@ -10461,16 +11401,19 @@ function meterFragment(presentCallCount, total) {
|
|
|
10461
11401
|
return { presentCallCount, total };
|
|
10462
11402
|
}
|
|
10463
11403
|
function accountingItemFromRow(row) {
|
|
10464
|
-
const usage =
|
|
11404
|
+
const usage = usageFromRow2(row, parseStringArray(row.requester_group_ids));
|
|
10465
11405
|
const cost = row.cost_status === null ? undefined : costFromJoinedRow(row);
|
|
10466
11406
|
const attribution = row.attribution_kind === null ? undefined : row.attribution_kind === "agent" ? {
|
|
10467
11407
|
kind: "agent",
|
|
10468
|
-
agentId: required2(row.attribution_agent_id),
|
|
10469
11408
|
agentRunId: required2(row.attribution_agent_run_id),
|
|
10470
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)
|
|
10471
11414
|
} : {
|
|
10472
11415
|
kind: "workflowAgent",
|
|
10473
|
-
|
|
11416
|
+
agentStepId: required2(row.attribution_agent_step_id),
|
|
10474
11417
|
nodeRunId: required2(row.attribution_node_run_id),
|
|
10475
11418
|
workflowId: required2(row.attribution_workflow_id),
|
|
10476
11419
|
workflowRunId: required2(row.attribution_workflow_run_id)
|
|
@@ -10481,7 +11424,7 @@ function accountingItemFromRow(row) {
|
|
|
10481
11424
|
...cost ? { cost } : {}
|
|
10482
11425
|
};
|
|
10483
11426
|
}
|
|
10484
|
-
function
|
|
11427
|
+
function usageFromRow2(row, requesterGroupIds) {
|
|
10485
11428
|
const number = (value) => value === null ? undefined : Number(value);
|
|
10486
11429
|
return {
|
|
10487
11430
|
id: row.id,
|
|
@@ -10491,7 +11434,11 @@ function usageFromRow(row, requesterGroupIds) {
|
|
|
10491
11434
|
callId: row.call_id,
|
|
10492
11435
|
requesterGroupIds,
|
|
10493
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
|
+
},
|
|
10494
11440
|
requestedModelId: row.requested_model_id,
|
|
11441
|
+
...row.requested_reasoning === null ? {} : { requestedReasoning: requestedReasoningFromRow(row.requested_reasoning) },
|
|
10495
11442
|
...row.response_model_id === null ? {} : { responseModelId: row.response_model_id },
|
|
10496
11443
|
responseId: row.response_id,
|
|
10497
11444
|
usage: {
|
|
@@ -10540,14 +11487,12 @@ function costFromJoinedRow(row) {
|
|
|
10540
11487
|
}
|
|
10541
11488
|
function costRecord(input) {
|
|
10542
11489
|
const details = parseAiModelCallCostDetails(input.details);
|
|
11490
|
+
const priceSource = details.priceSource === undefined ? undefined : { ...details.priceSource, observedAt: new Date(details.priceSource.observedAt) };
|
|
10543
11491
|
const base = {
|
|
11492
|
+
...details.estimate === undefined ? {} : { estimate: details.estimate },
|
|
10544
11493
|
projectId: input.projectId,
|
|
10545
11494
|
usageRecordId: input.usageId,
|
|
10546
11495
|
pricingContext: details.pricingContext,
|
|
10547
|
-
priceSource: {
|
|
10548
|
-
...details.priceSource,
|
|
10549
|
-
observedAt: new Date(details.priceSource.observedAt)
|
|
10550
|
-
},
|
|
10551
11496
|
ratedAt: new Date(input.ratedAt)
|
|
10552
11497
|
};
|
|
10553
11498
|
if (input.status === "rated") {
|
|
@@ -10555,6 +11500,7 @@ function costRecord(input) {
|
|
|
10555
11500
|
...base,
|
|
10556
11501
|
status: "rated",
|
|
10557
11502
|
billingIdentity: { providerId: required2(input.providerId), modelId: required2(input.modelId) },
|
|
11503
|
+
priceSource: required2(priceSource),
|
|
10558
11504
|
money: { currency: required2(input.currency), amountNanos: required2(input.amountNanos) },
|
|
10559
11505
|
components: required2(details.components)
|
|
10560
11506
|
});
|
|
@@ -10563,6 +11509,7 @@ function costRecord(input) {
|
|
|
10563
11509
|
...base,
|
|
10564
11510
|
status: "unpriceable",
|
|
10565
11511
|
...input.providerId && input.modelId ? { billingIdentity: { providerId: input.providerId, modelId: input.modelId } } : {},
|
|
11512
|
+
...priceSource === undefined ? {} : { priceSource },
|
|
10566
11513
|
reason: required2(input.reason),
|
|
10567
11514
|
...details.missingMeters ? { missingMeters: details.missingMeters } : {}
|
|
10568
11515
|
});
|
|
@@ -10643,222 +11590,611 @@ function required2(value) {
|
|
|
10643
11590
|
return value;
|
|
10644
11591
|
}
|
|
10645
11592
|
|
|
10646
|
-
// src/pg-ai-
|
|
11593
|
+
// src/pg-ai-limit-storage.ts
|
|
10647
11594
|
import {
|
|
10648
|
-
|
|
10649
|
-
|
|
10650
|
-
|
|
10651
|
-
|
|
10652
|
-
|
|
10653
|
-
|
|
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 {
|
|
10654
11618
|
sql;
|
|
11619
|
+
operations = new AiLimitOperationLock;
|
|
10655
11620
|
constructor(sql) {
|
|
10656
11621
|
this.sql = sql;
|
|
10657
11622
|
}
|
|
10658
|
-
|
|
10659
|
-
|
|
10660
|
-
|
|
10661
|
-
|
|
10662
|
-
|
|
10663
|
-
|
|
10664
|
-
|
|
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);
|
|
10665
11631
|
const inserted = await tx`
|
|
10666
|
-
INSERT INTO
|
|
10667
|
-
project_id,
|
|
10668
|
-
|
|
10669
|
-
execution_id,
|
|
10670
|
-
attempt,
|
|
10671
|
-
call_id,
|
|
10672
|
-
provider_id,
|
|
10673
|
-
requested_model_id,
|
|
10674
|
-
response_model_id,
|
|
10675
|
-
response_id,
|
|
10676
|
-
input_tokens,
|
|
10677
|
-
output_tokens,
|
|
10678
|
-
total_tokens,
|
|
10679
|
-
uncached_input_tokens,
|
|
10680
|
-
cache_read_input_tokens,
|
|
10681
|
-
cache_write_input_tokens,
|
|
10682
|
-
text_output_tokens,
|
|
10683
|
-
reasoning_output_tokens,
|
|
10684
|
-
reporting_status,
|
|
10685
|
-
raw_usage,
|
|
10686
|
-
occurred_at,
|
|
10687
|
-
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
|
|
10688
11635
|
) VALUES (
|
|
10689
|
-
${
|
|
10690
|
-
${
|
|
10691
|
-
${
|
|
10692
|
-
${record.attempt},
|
|
10693
|
-
${record.callId},
|
|
10694
|
-
${record.providerId},
|
|
10695
|
-
${record.requestedModelId},
|
|
10696
|
-
${record.responseModelId ?? null},
|
|
10697
|
-
${record.responseId},
|
|
10698
|
-
${record.usage.inputTokens ?? null},
|
|
10699
|
-
${record.usage.outputTokens ?? null},
|
|
10700
|
-
${record.usage.totalTokens ?? null},
|
|
10701
|
-
${record.usage.uncachedInputTokens ?? null},
|
|
10702
|
-
${record.usage.cacheReadInputTokens ?? null},
|
|
10703
|
-
${record.usage.cacheWriteInputTokens ?? null},
|
|
10704
|
-
${record.usage.textOutputTokens ?? null},
|
|
10705
|
-
${record.usage.reasoningOutputTokens ?? null},
|
|
10706
|
-
${record.usage.reportingStatus},
|
|
10707
|
-
${record.rawUsage === undefined ? null : JSON.stringify(record.rawUsage)}::text::jsonb,
|
|
10708
|
-
${record.occurredAt},
|
|
10709
|
-
${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}
|
|
10710
11639
|
)
|
|
10711
11640
|
ON CONFLICT DO NOTHING
|
|
10712
11641
|
RETURNING id
|
|
10713
11642
|
`;
|
|
10714
|
-
if (inserted.length === 0)
|
|
10715
|
-
|
|
10716
|
-
|
|
10717
|
-
|
|
10718
|
-
|
|
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 };
|
|
10719
11709
|
}
|
|
10720
|
-
|
|
10721
|
-
|
|
10722
|
-
|
|
10723
|
-
|
|
10724
|
-
|
|
10725
|
-
|
|
10726
|
-
|
|
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
|
|
10727
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
|
+
}
|
|
10728
11808
|
}
|
|
10729
|
-
return { record: await this.requireById(tx, record.projectId, record.id), created: true };
|
|
10730
11809
|
});
|
|
10731
11810
|
}
|
|
10732
|
-
async
|
|
10733
|
-
const
|
|
10734
|
-
|
|
10735
|
-
|
|
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);
|
|
10736
11836
|
});
|
|
10737
|
-
return summary;
|
|
10738
11837
|
}
|
|
10739
|
-
async
|
|
10740
|
-
|
|
10741
|
-
|
|
10742
|
-
|
|
10743
|
-
|
|
10744
|
-
|
|
10745
|
-
|
|
10746
|
-
|
|
10747
|
-
|
|
10748
|
-
|
|
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
|
+
});
|
|
10749
11854
|
}
|
|
10750
|
-
async
|
|
10751
|
-
|
|
10752
|
-
|
|
10753
|
-
|
|
10754
|
-
|
|
10755
|
-
|
|
10756
|
-
|
|
10757
|
-
|
|
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
|
+
});
|
|
10758
11863
|
}
|
|
10759
|
-
|
|
10760
|
-
|
|
10761
|
-
SELECT * FROM
|
|
10762
|
-
WHERE project_id = ${
|
|
10763
|
-
|
|
10764
|
-
AND attempt = ${record.attempt}
|
|
10765
|
-
AND call_id = ${record.callId}
|
|
10766
|
-
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
|
|
10767
11869
|
`;
|
|
10768
|
-
return rows[0] ? this.rowToRecord(sql, rows[0]) : null;
|
|
10769
11870
|
}
|
|
10770
|
-
async
|
|
11871
|
+
async requirePolicy(sql, projectId, id) {
|
|
10771
11872
|
const [row] = await sql`
|
|
10772
|
-
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}
|
|
10773
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
|
+
}
|
|
10774
11912
|
if (!row)
|
|
10775
|
-
throw new Error(
|
|
10776
|
-
return
|
|
11913
|
+
throw new Error("[SixbPg] Failed to initialize AI limit period state.");
|
|
11914
|
+
return policyStatus(policy, period, row, existingGroups);
|
|
10777
11915
|
}
|
|
10778
|
-
async
|
|
10779
|
-
const
|
|
10780
|
-
|
|
10781
|
-
|
|
10782
|
-
|
|
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
|
|
10783
11951
|
`;
|
|
10784
|
-
return {
|
|
10785
|
-
id: row.id,
|
|
10786
|
-
projectId: row.project_id,
|
|
10787
|
-
executionId: row.execution_id,
|
|
10788
|
-
attempt: Number(row.attempt),
|
|
10789
|
-
callId: row.call_id,
|
|
10790
|
-
requesterGroupIds: groupRows.map((group) => group.group_id),
|
|
10791
|
-
providerId: row.provider_id,
|
|
10792
|
-
requestedModelId: row.requested_model_id,
|
|
10793
|
-
...row.response_model_id === null ? {} : { responseModelId: row.response_model_id },
|
|
10794
|
-
responseId: row.response_id,
|
|
10795
|
-
usage: {
|
|
10796
|
-
...usageFromRow2(row),
|
|
10797
|
-
...row.total_tokens === null ? {} : { totalTokens: Number(row.total_tokens) },
|
|
10798
|
-
reportingStatus: row.reporting_status
|
|
10799
|
-
},
|
|
10800
|
-
...row.raw_usage === null ? {} : { rawUsage: rawUsageFromRow(row.raw_usage) },
|
|
10801
|
-
occurredAt: new Date(row.occurred_at),
|
|
10802
|
-
recordedAt: new Date(row.recorded_at)
|
|
10803
|
-
};
|
|
10804
11952
|
}
|
|
10805
|
-
|
|
10806
|
-
const
|
|
10807
|
-
|
|
10808
|
-
|
|
10809
|
-
|
|
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}
|
|
10810
11984
|
)
|
|
10811
|
-
SELECT usage.* FROM requested
|
|
10812
|
-
JOIN ai_model_call_usage AS usage
|
|
10813
|
-
ON usage.project_id = ${input.projectId}
|
|
10814
|
-
AND usage.execution_id = requested.execution_id
|
|
10815
11985
|
`;
|
|
10816
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
|
+
}
|
|
10817
12078
|
async assertExecutionExists(sql, projectId, executionId) {
|
|
10818
|
-
const [
|
|
12079
|
+
const [row] = await sql`
|
|
10819
12080
|
SELECT EXISTS (
|
|
10820
12081
|
SELECT 1 FROM executions WHERE project_id = ${projectId} AND id = ${executionId}
|
|
10821
12082
|
) AS exists
|
|
10822
12083
|
`;
|
|
10823
|
-
if (
|
|
12084
|
+
if (row?.exists)
|
|
10824
12085
|
return;
|
|
10825
|
-
throw new
|
|
12086
|
+
throw new AiLimitStorageError("missing_execution", `[SixbPg] AI limit execution '${executionId}' does not exist in project '${projectId}'.`);
|
|
10826
12087
|
}
|
|
10827
12088
|
}
|
|
10828
|
-
function
|
|
10829
|
-
|
|
10830
|
-
|
|
10831
|
-
|
|
10832
|
-
|
|
10833
|
-
|
|
10834
|
-
|
|
10835
|
-
|
|
10836
|
-
|
|
10837
|
-
|
|
10838
|
-
|
|
10839
|
-
|
|
10840
|
-
|
|
10841
|
-
|
|
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
|
+
};
|
|
10842
12104
|
}
|
|
10843
|
-
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 });
|
|
10844
12112
|
return {
|
|
10845
|
-
|
|
10846
|
-
|
|
10847
|
-
|
|
10848
|
-
|
|
10849
|
-
|
|
10850
|
-
|
|
10851
|
-
|
|
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)
|
|
10852
12124
|
};
|
|
10853
12125
|
}
|
|
10854
|
-
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) {
|
|
10855
12175
|
return typeof value === "string" ? JSON.parse(value) : value;
|
|
10856
12176
|
}
|
|
10857
|
-
function
|
|
10858
|
-
if (
|
|
10859
|
-
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.");
|
|
10860
12180
|
}
|
|
10861
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
|
+
}
|
|
10862
12198
|
|
|
10863
12199
|
// src/pg-client.ts
|
|
10864
12200
|
import postgres from "postgres";
|
|
@@ -14712,15 +16048,15 @@ class PgWorkflowAgentNodeRunStorage {
|
|
|
14712
16048
|
projectId: input.projectId,
|
|
14713
16049
|
executionId: input.executionId,
|
|
14714
16050
|
nodeRunId: input.nodeRunId,
|
|
14715
|
-
|
|
16051
|
+
actorId: input.actorId,
|
|
14716
16052
|
workflowExecutionId: parent.execution_id
|
|
14717
16053
|
});
|
|
14718
16054
|
try {
|
|
14719
16055
|
const [row] = await this.sql`
|
|
14720
16056
|
INSERT INTO workflow_agent_node_runs (
|
|
14721
|
-
project_id, node_run_id, execution_id,
|
|
16057
|
+
project_id, node_run_id, execution_id, actor_id, status, prompt, attempt, created_at
|
|
14722
16058
|
) VALUES (
|
|
14723
|
-
${input.projectId}, ${input.nodeRunId}, ${input.executionId}, ${input.
|
|
16059
|
+
${input.projectId}, ${input.nodeRunId}, ${input.executionId}, ${input.actorId},
|
|
14724
16060
|
${"queued"}, ${input.prompt}, ${0}, ${input.createdAt ?? new Date}
|
|
14725
16061
|
) RETURNING *
|
|
14726
16062
|
`;
|
|
@@ -14837,9 +16173,9 @@ class PgWorkflowAgentNodeRunStorage {
|
|
|
14837
16173
|
const whereClauses = ["project_id = $1"];
|
|
14838
16174
|
const params = [input.projectId];
|
|
14839
16175
|
let index = 2;
|
|
14840
|
-
if (input.
|
|
14841
|
-
whereClauses.push(`
|
|
14842
|
-
params.push(input.
|
|
16176
|
+
if (input.actorId) {
|
|
16177
|
+
whereClauses.push(`actor_id = $${index++}`);
|
|
16178
|
+
params.push(input.actorId);
|
|
14843
16179
|
}
|
|
14844
16180
|
index = appendRunListFilters(whereClauses, params, index, input);
|
|
14845
16181
|
const result = await queryRunList({
|
|
@@ -15116,7 +16452,7 @@ function rowToWorkflowAgentNodeRunRecord(row) {
|
|
|
15116
16452
|
projectId: row.project_id,
|
|
15117
16453
|
nodeRunId: row.node_run_id,
|
|
15118
16454
|
executionId: row.execution_id,
|
|
15119
|
-
|
|
16455
|
+
actorId: row.actor_id,
|
|
15120
16456
|
status: row.status,
|
|
15121
16457
|
prompt: row.prompt,
|
|
15122
16458
|
...row.model_id ? { modelId: row.model_id } : {},
|
|
@@ -15181,6 +16517,7 @@ class PostgresStorage {
|
|
|
15181
16517
|
agents;
|
|
15182
16518
|
aiUsage;
|
|
15183
16519
|
aiCosts;
|
|
16520
|
+
aiLimits;
|
|
15184
16521
|
actionRuns;
|
|
15185
16522
|
pipelineRuns;
|
|
15186
16523
|
workflowRuns;
|
|
@@ -15231,6 +16568,7 @@ class PostgresStorage {
|
|
|
15231
16568
|
this.agents = createAgentOperationScope(stores.agents, scope);
|
|
15232
16569
|
this.aiUsage = createOperationScopedFacade(stores.aiUsage, scope);
|
|
15233
16570
|
this.aiCosts = createOperationScopedFacade(stores.aiCosts, scope);
|
|
16571
|
+
this.aiLimits = createOperationScopedFacade(stores.aiLimits, scope);
|
|
15234
16572
|
this.actionRuns = createOperationScopedFacade(stores.actionRuns, scope);
|
|
15235
16573
|
this.pipelineRuns = createOperationScopedFacade(stores.pipelineRuns, scope);
|
|
15236
16574
|
this.workflowRuns = createWorkflowRunOperationScope(stores.workflowRuns, scope);
|
|
@@ -15266,6 +16604,7 @@ class PostgresStorage {
|
|
|
15266
16604
|
transactionContext.materializations.assertCommittable();
|
|
15267
16605
|
return result;
|
|
15268
16606
|
} finally {
|
|
16607
|
+
await txStorage.aiLimits.settleOperations();
|
|
15269
16608
|
transactionContext.active = false;
|
|
15270
16609
|
txStorage.ontology.deactivateSessions();
|
|
15271
16610
|
transactionContext.materializations.deactivate();
|
|
@@ -15320,6 +16659,7 @@ function createPostgresStores(sql, options) {
|
|
|
15320
16659
|
agents: new PgAgentStorage({ sql, executions }),
|
|
15321
16660
|
aiUsage: new PgAiUsageStorage(sql),
|
|
15322
16661
|
aiCosts: new PgAiCostStorage(sql),
|
|
16662
|
+
aiLimits: new PgAiLimitStorage(sql),
|
|
15323
16663
|
actionRuns: new PgActionRunStorage(sql, executions),
|
|
15324
16664
|
pipelineRuns: new PgPipelineRunStorage(sql, executions),
|
|
15325
16665
|
workflowRuns: new PgWorkflowRunStorage(sql, executions),
|
|
@@ -15346,4 +16686,4 @@ export {
|
|
|
15346
16686
|
PostgresStorage
|
|
15347
16687
|
};
|
|
15348
16688
|
|
|
15349
|
-
//# debugId=
|
|
16689
|
+
//# debugId=A5AA822A44042C7A64756E2164756E21
|