mirror of
https://github.com/whyour/qinglong.git
synced 2026-09-21 09:58:46 +08:00
341 lines
11 KiB
TypeScript
341 lines
11 KiB
TypeScript
import type { PostgresClient, PostgresPool } from '@qinglong/runtime-core';
|
|
import {
|
|
InvalidToolInvocationArtifactError,
|
|
MAX_TOOL_INVOCATION_INPUT_ARTIFACT_JSON_BYTES,
|
|
MAX_TOOL_INVOCATION_PREVIEW_ARTIFACT_JSON_BYTES,
|
|
ToolInvocationArtifactConflictError,
|
|
ToolInvocationArtifactUnavailableError,
|
|
normalizeToolInvocationInputArtifact,
|
|
normalizeToolInvocationPreviewArtifact,
|
|
type ToolInvocationArtifactRepository,
|
|
type ToolInvocationInputArtifact,
|
|
type ToolInvocationPreviewArtifact,
|
|
} from '@qinglong/runtime-core/tool-invocation-artifact';
|
|
|
|
import {
|
|
POSTGRES_DEFINITION_RETRYABLE_SQL_STATES,
|
|
POSTGRES_DEFINITION_TRANSACTION_ATTEMPTS,
|
|
configurePostgresDefinitionTransaction,
|
|
postgresRequiredInteger,
|
|
postgresRequiredJsonObject,
|
|
postgresRequiredString,
|
|
postgresSqlState,
|
|
rollbackPostgresDefinitionTransaction,
|
|
} from '../repository/definitionRepositorySupport';
|
|
|
|
type Row = Record<string, unknown>;
|
|
type Queryable = Pick<PostgresPool, 'query'> | Pick<PostgresClient, 'query'>;
|
|
|
|
const ARTIFACT_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/;
|
|
|
|
function unavailable(cause?: unknown): ToolInvocationArtifactUnavailableError {
|
|
return new ToolInvocationArtifactUnavailableError({
|
|
cause: cause instanceof Error ? cause : undefined,
|
|
});
|
|
}
|
|
|
|
function artifactId(value: unknown): string {
|
|
if (typeof value !== 'string' || !ARTIFACT_ID_PATTERN.test(value)) {
|
|
throw new InvalidToolInvocationArtifactError('artifact id is invalid');
|
|
}
|
|
return value;
|
|
}
|
|
|
|
function serialize(value: unknown, maximumBytes: number): string {
|
|
const json = JSON.stringify(value);
|
|
if (Buffer.byteLength(json, 'utf8') > maximumBytes) {
|
|
throw new InvalidToolInvocationArtifactError(
|
|
'durable Artifact JSON exceeds its budget',
|
|
);
|
|
}
|
|
return json;
|
|
}
|
|
|
|
function mappedError(error: unknown): Error {
|
|
if (
|
|
error instanceof InvalidToolInvocationArtifactError ||
|
|
error instanceof ToolInvocationArtifactConflictError ||
|
|
error instanceof ToolInvocationArtifactUnavailableError
|
|
) {
|
|
return error;
|
|
}
|
|
const state = postgresSqlState(error);
|
|
if (state === '23503' || state === '23505' || state === '23514') {
|
|
return new ToolInvocationArtifactConflictError();
|
|
}
|
|
return unavailable(error);
|
|
}
|
|
|
|
function requiredString(value: unknown): string {
|
|
return postgresRequiredString(value, unavailable);
|
|
}
|
|
|
|
function requiredInteger(value: unknown): number {
|
|
return postgresRequiredInteger(value, unavailable);
|
|
}
|
|
|
|
export class PostgresToolInvocationArtifactRepository
|
|
implements ToolInvocationArtifactRepository
|
|
{
|
|
constructor(private readonly pool: PostgresPool) {
|
|
if (
|
|
!pool ||
|
|
typeof pool.query !== 'function' ||
|
|
typeof pool.connect !== 'function'
|
|
) {
|
|
throw new TypeError(
|
|
'PostgreSQL Tool invocation Artifact repository options are invalid',
|
|
);
|
|
}
|
|
}
|
|
|
|
private parseInput(row: Row): Readonly<ToolInvocationInputArtifact> {
|
|
try {
|
|
const artifact = normalizeToolInvocationInputArtifact(
|
|
postgresRequiredJsonObject(
|
|
row.artifactJson,
|
|
unavailable,
|
|
) as unknown as ToolInvocationInputArtifact,
|
|
);
|
|
if (
|
|
artifact.artifactId !== requiredString(row.artifactId) ||
|
|
artifact.projectId !== requiredString(row.projectId) ||
|
|
artifact.actionRef !== requiredString(row.actionRef) ||
|
|
artifact.inputDigest !== requiredString(row.inputDigest) ||
|
|
artifact.invocationActionDigest !==
|
|
requiredString(row.invocationActionDigest) ||
|
|
artifact.artifactDigest !== requiredString(row.artifactDigest) ||
|
|
artifact.keyId !== requiredString(row.keyId) ||
|
|
artifact.algorithm !== requiredString(row.algorithm) ||
|
|
artifact.plaintextBytes !== requiredInteger(row.plaintextBytes) ||
|
|
artifact.sealedAtMs !== requiredInteger(row.sealedAtMs) ||
|
|
Buffer.byteLength(JSON.stringify(artifact), 'utf8') >
|
|
MAX_TOOL_INVOCATION_INPUT_ARTIFACT_JSON_BYTES
|
|
) {
|
|
throw unavailable();
|
|
}
|
|
return artifact;
|
|
} catch (error) {
|
|
if (error instanceof ToolInvocationArtifactUnavailableError) throw error;
|
|
throw unavailable();
|
|
}
|
|
}
|
|
|
|
private parsePreview(
|
|
row: Row,
|
|
): Readonly<ToolInvocationPreviewArtifact> {
|
|
try {
|
|
const artifact = normalizeToolInvocationPreviewArtifact(
|
|
postgresRequiredJsonObject(
|
|
row.artifactJson,
|
|
unavailable,
|
|
) as unknown as ToolInvocationPreviewArtifact,
|
|
);
|
|
if (
|
|
artifact.artifactId !== requiredString(row.artifactId) ||
|
|
artifact.projectId !== requiredString(row.projectId) ||
|
|
artifact.actionRef !== requiredString(row.actionRef) ||
|
|
artifact.actionDigest !== requiredString(row.actionDigest) ||
|
|
artifact.previewDigest !== requiredString(row.previewDigest) ||
|
|
artifact.redactionContractDigest !==
|
|
requiredString(row.redactionContractDigest) ||
|
|
artifact.artifactDigest !== requiredString(row.artifactDigest) ||
|
|
artifact.byteLength !== requiredInteger(row.byteLength) ||
|
|
artifact.sealedAtMs !== requiredInteger(row.sealedAtMs) ||
|
|
Buffer.byteLength(JSON.stringify(artifact), 'utf8') >
|
|
MAX_TOOL_INVOCATION_PREVIEW_ARTIFACT_JSON_BYTES
|
|
) {
|
|
throw unavailable();
|
|
}
|
|
return artifact;
|
|
} catch (error) {
|
|
if (error instanceof ToolInvocationArtifactUnavailableError) throw error;
|
|
throw unavailable();
|
|
}
|
|
}
|
|
|
|
private async storedInput(
|
|
queryable: Queryable,
|
|
id: string,
|
|
): Promise<Readonly<ToolInvocationInputArtifact> | null> {
|
|
const result = await queryable.query<Row>(
|
|
`SELECT
|
|
artifact_id AS "artifactId",
|
|
project_id AS "projectId",
|
|
action_ref AS "actionRef",
|
|
input_digest AS "inputDigest",
|
|
invocation_action_digest AS "invocationActionDigest",
|
|
artifact_digest AS "artifactDigest",
|
|
key_id AS "keyId",
|
|
algorithm,
|
|
plaintext_bytes AS "plaintextBytes",
|
|
sealed_at_ms AS "sealedAtMs",
|
|
artifact_json AS "artifactJson"
|
|
FROM "ql3"."tool_invocation_input_artifacts"
|
|
WHERE artifact_id = $1
|
|
LIMIT 2`,
|
|
[id],
|
|
);
|
|
if (result.rows.length > 1) throw unavailable();
|
|
return result.rows[0] ? this.parseInput(result.rows[0]) : null;
|
|
}
|
|
|
|
private async storedPreview(
|
|
queryable: Queryable,
|
|
id: string,
|
|
): Promise<Readonly<ToolInvocationPreviewArtifact> | null> {
|
|
const result = await queryable.query<Row>(
|
|
`SELECT
|
|
artifact_id AS "artifactId",
|
|
project_id AS "projectId",
|
|
action_ref AS "actionRef",
|
|
action_digest AS "actionDigest",
|
|
preview_digest AS "previewDigest",
|
|
redaction_contract_digest AS "redactionContractDigest",
|
|
artifact_digest AS "artifactDigest",
|
|
byte_length AS "byteLength",
|
|
sealed_at_ms AS "sealedAtMs",
|
|
artifact_json AS "artifactJson"
|
|
FROM "ql3"."tool_invocation_preview_artifacts"
|
|
WHERE artifact_id = $1
|
|
LIMIT 2`,
|
|
[id],
|
|
);
|
|
if (result.rows.length > 1) throw unavailable();
|
|
return result.rows[0] ? this.parsePreview(result.rows[0]) : null;
|
|
}
|
|
|
|
async put(
|
|
inputValue: ToolInvocationInputArtifact,
|
|
previewValue: ToolInvocationPreviewArtifact,
|
|
): Promise<Readonly<{ status: 'inserted' | 'existing' }>> {
|
|
const input = normalizeToolInvocationInputArtifact(inputValue);
|
|
const preview = normalizeToolInvocationPreviewArtifact(previewValue);
|
|
if (
|
|
input.projectId !== preview.projectId ||
|
|
input.actionRef !== preview.actionRef ||
|
|
input.sealedAtMs !== preview.sealedAtMs
|
|
) {
|
|
throw new ToolInvocationArtifactConflictError();
|
|
}
|
|
serialize(input, MAX_TOOL_INVOCATION_INPUT_ARTIFACT_JSON_BYTES);
|
|
serialize(preview, MAX_TOOL_INVOCATION_PREVIEW_ARTIFACT_JSON_BYTES);
|
|
for (
|
|
let attempt = 1;
|
|
attempt <= POSTGRES_DEFINITION_TRANSACTION_ATTEMPTS;
|
|
attempt += 1
|
|
) {
|
|
let client: PostgresClient;
|
|
try {
|
|
client = await this.pool.connect();
|
|
} catch {
|
|
throw unavailable();
|
|
}
|
|
let began = false;
|
|
try {
|
|
await configurePostgresDefinitionTransaction(client);
|
|
began = true;
|
|
const storedInput = await this.storedInput(client, input.artifactId);
|
|
const storedPreview = await this.storedPreview(
|
|
client,
|
|
preview.artifactId,
|
|
);
|
|
if (storedInput || storedPreview) {
|
|
if (
|
|
!storedInput ||
|
|
!storedPreview ||
|
|
JSON.stringify(storedInput) !== JSON.stringify(input) ||
|
|
JSON.stringify(storedPreview) !== JSON.stringify(preview)
|
|
) {
|
|
throw new ToolInvocationArtifactConflictError();
|
|
}
|
|
await client.query('COMMIT');
|
|
began = false;
|
|
return Object.freeze({ status: 'existing' });
|
|
}
|
|
await client.query(
|
|
`INSERT INTO "ql3"."tool_invocation_input_artifacts" (
|
|
artifact_id, project_id, action_ref, input_digest,
|
|
invocation_action_digest, artifact_digest, key_id, algorithm,
|
|
plaintext_bytes, sealed_at_ms, artifact_json
|
|
) VALUES (
|
|
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::jsonb
|
|
)`,
|
|
[
|
|
input.artifactId,
|
|
input.projectId,
|
|
input.actionRef,
|
|
input.inputDigest,
|
|
input.invocationActionDigest,
|
|
input.artifactDigest,
|
|
input.keyId,
|
|
input.algorithm,
|
|
input.plaintextBytes,
|
|
input.sealedAtMs,
|
|
JSON.stringify(input),
|
|
],
|
|
);
|
|
await client.query(
|
|
`INSERT INTO "ql3"."tool_invocation_preview_artifacts" (
|
|
artifact_id, project_id, action_ref, action_digest,
|
|
preview_digest, redaction_contract_digest, artifact_digest,
|
|
byte_length, sealed_at_ms, artifact_json
|
|
) VALUES (
|
|
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::jsonb
|
|
)`,
|
|
[
|
|
preview.artifactId,
|
|
preview.projectId,
|
|
preview.actionRef,
|
|
preview.actionDigest,
|
|
preview.previewDigest,
|
|
preview.redactionContractDigest,
|
|
preview.artifactDigest,
|
|
preview.byteLength,
|
|
preview.sealedAtMs,
|
|
JSON.stringify(preview),
|
|
],
|
|
);
|
|
await client.query('COMMIT');
|
|
began = false;
|
|
return Object.freeze({ status: 'inserted' });
|
|
} catch (error) {
|
|
if (began) await rollbackPostgresDefinitionTransaction(client);
|
|
const state = postgresSqlState(error);
|
|
if (
|
|
attempt < POSTGRES_DEFINITION_TRANSACTION_ATTEMPTS &&
|
|
state !== undefined &&
|
|
(POSTGRES_DEFINITION_RETRYABLE_SQL_STATES.has(state) ||
|
|
state === '23505')
|
|
) {
|
|
continue;
|
|
}
|
|
throw mappedError(error);
|
|
} finally {
|
|
client.release();
|
|
}
|
|
}
|
|
throw unavailable();
|
|
}
|
|
|
|
async findInput(
|
|
value: string,
|
|
): Promise<Readonly<ToolInvocationInputArtifact> | null> {
|
|
try {
|
|
return await this.storedInput(this.pool, artifactId(value));
|
|
} catch (error) {
|
|
throw mappedError(error);
|
|
}
|
|
}
|
|
|
|
async findPreview(
|
|
value: string,
|
|
): Promise<Readonly<ToolInvocationPreviewArtifact> | null> {
|
|
try {
|
|
return await this.storedPreview(this.pool, artifactId(value));
|
|
} catch (error) {
|
|
throw mappedError(error);
|
|
}
|
|
}
|
|
}
|