postgresai 0.16.0-dev.1 → 0.16.0-dev.3
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/bin/postgres-ai.ts +728 -6
- package/dist/bin/postgres-ai.js +2496 -255
- package/dist/sql/06.helpers.sql +0 -122
- package/dist/sql/sql/06.helpers.sql +0 -122
- package/lib/dblab.ts +457 -0
- package/lib/init.ts +0 -8
- package/lib/joe.ts +730 -0
- package/lib/mcp-server.ts +597 -0
- package/lib/supabase.ts +0 -18
- package/package.json +1 -1
- package/sql/06.helpers.sql +0 -122
- package/test/dblab.cli.test.ts +339 -0
- package/test/dblab.test.ts +408 -0
- package/test/init.integration.test.ts +9 -79
- package/test/joe.cli.test.ts +298 -0
- package/test/joe.test.ts +726 -0
- package/test/monitoring.test.ts +54 -3
package/lib/joe.ts
ADDED
|
@@ -0,0 +1,730 @@
|
|
|
1
|
+
import * as fs from "fs";
|
|
2
|
+
import * as path from "path";
|
|
3
|
+
import * as crypto from "node:crypto";
|
|
4
|
+
import { formatHttpError, maskSecret, normalizeBaseUrl } from "./util";
|
|
5
|
+
import * as config from "./config";
|
|
6
|
+
|
|
7
|
+
/**
|
|
8
|
+
* Joe API v2 client (`postgres-ai` CLI / MCP surface).
|
|
9
|
+
*
|
|
10
|
+
* Every Joe command is a thin wrapper over the platform rpc `v1.joe_command_submit`
|
|
11
|
+
* (with the command's fixed `command`) plus a *client-side* status/result poll — there
|
|
12
|
+
* is no DB-side sync wrapper (a single PostgREST transaction cannot both submit a
|
|
13
|
+
* command and observe Joe's cross-connection callback). See SPEC §1/§5/§6.
|
|
14
|
+
*
|
|
15
|
+
* The backend rpcs (`joe_command_submit` / `joe_command_status` / `joe_command_result`
|
|
16
|
+
* and the org-level projects listing) do not exist yet — this module is written against
|
|
17
|
+
* the documented contract and exercised entirely with mocked fetch responses.
|
|
18
|
+
*/
|
|
19
|
+
|
|
20
|
+
/** The Joe command set (SPEC §5/§6 + the per-command contract matrix). */
|
|
21
|
+
export const JOE_COMMANDS = [
|
|
22
|
+
"plan",
|
|
23
|
+
"explain",
|
|
24
|
+
"exec",
|
|
25
|
+
"hypo",
|
|
26
|
+
"activity",
|
|
27
|
+
"terminate",
|
|
28
|
+
"reset",
|
|
29
|
+
"describe",
|
|
30
|
+
] as const;
|
|
31
|
+
|
|
32
|
+
export type JoeCommand = (typeof JOE_COMMANDS)[number];
|
|
33
|
+
|
|
34
|
+
/** Terminal + non-terminal command lifecycle states returned by the platform. */
|
|
35
|
+
export type JoeStatus = "queued" | "running" | "done" | "error" | "timed_out";
|
|
36
|
+
|
|
37
|
+
const TERMINAL_STATES: ReadonlySet<JoeStatus> = new Set<JoeStatus>([
|
|
38
|
+
"done",
|
|
39
|
+
"error",
|
|
40
|
+
"timed_out",
|
|
41
|
+
]);
|
|
42
|
+
|
|
43
|
+
/** Commands whose result reaches the caller's LLM — gated on the org `ai_enabled`. */
|
|
44
|
+
export const AI_ENABLED_COMMANDS: ReadonlySet<JoeCommand> = new Set<JoeCommand>([
|
|
45
|
+
"plan",
|
|
46
|
+
"explain",
|
|
47
|
+
"exec",
|
|
48
|
+
"hypo",
|
|
49
|
+
"activity",
|
|
50
|
+
"describe",
|
|
51
|
+
]);
|
|
52
|
+
|
|
53
|
+
/** Commands that execute against the disposable clone (at-most-once posture). */
|
|
54
|
+
const EXECUTING_COMMANDS: ReadonlySet<JoeCommand> = new Set<JoeCommand>([
|
|
55
|
+
"exec",
|
|
56
|
+
"explain",
|
|
57
|
+
]);
|
|
58
|
+
|
|
59
|
+
/** Default one-shot poll budget (SPEC §6: ≤ 25 s, then resume by handle). */
|
|
60
|
+
export const DEFAULT_BUDGET_MS = 25_000;
|
|
61
|
+
const DEFAULT_POLL_INTERVAL_MS = 800;
|
|
62
|
+
|
|
63
|
+
// ---------------------------------------------------------------------------
|
|
64
|
+
// Response shapes (documented contract — mocked in tests)
|
|
65
|
+
// ---------------------------------------------------------------------------
|
|
66
|
+
|
|
67
|
+
export interface JoeSubmitResponse {
|
|
68
|
+
command_id: string;
|
|
69
|
+
session_id: string | null;
|
|
70
|
+
status: JoeStatus;
|
|
71
|
+
}
|
|
72
|
+
|
|
73
|
+
export interface JoeStatusResponse {
|
|
74
|
+
command_id: string;
|
|
75
|
+
status: JoeStatus;
|
|
76
|
+
enqueued_at?: string | null;
|
|
77
|
+
started_at?: string | null;
|
|
78
|
+
finished_at?: string | null;
|
|
79
|
+
error?: string | null;
|
|
80
|
+
}
|
|
81
|
+
|
|
82
|
+
/**
|
|
83
|
+
* Command-appropriate result body (SPEC §1 + per-command contract matrix Part B).
|
|
84
|
+
* A single interface is used with per-command optional keys; only the keys for the
|
|
85
|
+
* submitted command are populated.
|
|
86
|
+
*/
|
|
87
|
+
export interface JoeCommandResult {
|
|
88
|
+
command_id: string;
|
|
89
|
+
command?: JoeCommand;
|
|
90
|
+
status: JoeStatus;
|
|
91
|
+
error: string | null;
|
|
92
|
+
// plan / explain
|
|
93
|
+
queryid?: string | null;
|
|
94
|
+
plan_fingerprint?: string | null;
|
|
95
|
+
plan_text?: string | null;
|
|
96
|
+
plan_json?: unknown;
|
|
97
|
+
// exec
|
|
98
|
+
row_count?: number | null;
|
|
99
|
+
result_rows?: unknown[] | null;
|
|
100
|
+
notices?: string[] | null;
|
|
101
|
+
// hypo
|
|
102
|
+
hypo_plan?: unknown;
|
|
103
|
+
hypo_used?: boolean | null;
|
|
104
|
+
// activity / describe
|
|
105
|
+
snapshot?: unknown;
|
|
106
|
+
// terminate
|
|
107
|
+
terminated?: boolean | null;
|
|
108
|
+
pid?: number | null;
|
|
109
|
+
// reset
|
|
110
|
+
reset?: boolean | null;
|
|
111
|
+
}
|
|
112
|
+
|
|
113
|
+
export interface ProjectListItem {
|
|
114
|
+
project_id: number;
|
|
115
|
+
alias: string | null;
|
|
116
|
+
name: string | null;
|
|
117
|
+
label: string | null;
|
|
118
|
+
/** Whether the project's single Joe instance is `joe_api_v2_enabled` + conformant. */
|
|
119
|
+
joe_ready: boolean;
|
|
120
|
+
/** Whether the project's DBLab tunnel is connected. */
|
|
121
|
+
tunnel: boolean;
|
|
122
|
+
}
|
|
123
|
+
|
|
124
|
+
// ---------------------------------------------------------------------------
|
|
125
|
+
// Low-level rpc caller
|
|
126
|
+
// ---------------------------------------------------------------------------
|
|
127
|
+
|
|
128
|
+
interface RpcCallParams {
|
|
129
|
+
apiKey: string;
|
|
130
|
+
apiBaseUrl: string;
|
|
131
|
+
fn: string;
|
|
132
|
+
body: Record<string, unknown>;
|
|
133
|
+
operation: string;
|
|
134
|
+
debug?: boolean;
|
|
135
|
+
}
|
|
136
|
+
|
|
137
|
+
async function callRpc<T>(params: RpcCallParams): Promise<T> {
|
|
138
|
+
const { apiKey, apiBaseUrl, fn, body, operation, debug } = params;
|
|
139
|
+
if (!apiKey) {
|
|
140
|
+
throw new Error("API key is required");
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
const base = normalizeBaseUrl(apiBaseUrl);
|
|
144
|
+
const url = new URL(`${base}/rpc/${fn}`);
|
|
145
|
+
const payload = JSON.stringify(body);
|
|
146
|
+
|
|
147
|
+
const headers: Record<string, string> = {
|
|
148
|
+
"access-token": apiKey,
|
|
149
|
+
"Content-Type": "application/json",
|
|
150
|
+
"Connection": "close",
|
|
151
|
+
};
|
|
152
|
+
|
|
153
|
+
if (debug) {
|
|
154
|
+
const debugHeaders: Record<string, string> = { ...headers, "access-token": maskSecret(apiKey) };
|
|
155
|
+
console.error(`Debug: POST URL: ${url.toString()}`);
|
|
156
|
+
console.error(`Debug: Request headers: ${JSON.stringify(debugHeaders)}`);
|
|
157
|
+
console.error(`Debug: Request body: ${payload}`);
|
|
158
|
+
}
|
|
159
|
+
|
|
160
|
+
const response = await fetch(url.toString(), {
|
|
161
|
+
method: "POST",
|
|
162
|
+
headers,
|
|
163
|
+
body: payload,
|
|
164
|
+
});
|
|
165
|
+
|
|
166
|
+
const text = await response.text();
|
|
167
|
+
|
|
168
|
+
if (debug) {
|
|
169
|
+
console.error(`Debug: Response status: ${response.status}`);
|
|
170
|
+
console.error(`Debug: Response body: ${text}`);
|
|
171
|
+
}
|
|
172
|
+
|
|
173
|
+
if (!response.ok) {
|
|
174
|
+
// PostgREST maps a custom `PTxyz` sqlstate to HTTP status `xyz`, so PT403 →
|
|
175
|
+
// HTTP 403, PT404 → 404, PT429 → 429, etc. formatHttpError surfaces the JSON
|
|
176
|
+
// `message` (e.g. the "AI features disabled for this org" text for AC36).
|
|
177
|
+
throw new Error(formatHttpError(operation, response.status, text));
|
|
178
|
+
}
|
|
179
|
+
|
|
180
|
+
try {
|
|
181
|
+
return JSON.parse(text) as T;
|
|
182
|
+
} catch {
|
|
183
|
+
throw new Error(`${operation}: failed to parse response: ${text}`);
|
|
184
|
+
}
|
|
185
|
+
}
|
|
186
|
+
|
|
187
|
+
// ---------------------------------------------------------------------------
|
|
188
|
+
// Individual rpc client functions
|
|
189
|
+
// ---------------------------------------------------------------------------
|
|
190
|
+
|
|
191
|
+
export interface SubmitCommandParams {
|
|
192
|
+
apiKey: string;
|
|
193
|
+
apiBaseUrl: string;
|
|
194
|
+
command: JoeCommand;
|
|
195
|
+
projectId: number;
|
|
196
|
+
sql?: string | null;
|
|
197
|
+
args?: Record<string, unknown> | null;
|
|
198
|
+
/** 64-bit session id, carried verbatim as a string to avoid precision loss. */
|
|
199
|
+
sessionId?: string | null;
|
|
200
|
+
idempotencyKey?: string | null;
|
|
201
|
+
debug?: boolean;
|
|
202
|
+
}
|
|
203
|
+
|
|
204
|
+
/** Submit a Joe command (`v1.joe_command_submit`); returns the queued handle. */
|
|
205
|
+
export async function submitCommand(params: SubmitCommandParams): Promise<JoeSubmitResponse> {
|
|
206
|
+
const { apiKey, apiBaseUrl, command, projectId, sql, args, sessionId, idempotencyKey, debug } = params;
|
|
207
|
+
if (!JOE_COMMANDS.includes(command)) {
|
|
208
|
+
throw new Error(`Unknown Joe command: ${command}`);
|
|
209
|
+
}
|
|
210
|
+
const body: Record<string, unknown> = {
|
|
211
|
+
project_id: projectId,
|
|
212
|
+
command,
|
|
213
|
+
sql: sql ?? null,
|
|
214
|
+
args: args ?? null,
|
|
215
|
+
// session_id is a bigint server-side; sent as a numeric string when known so a
|
|
216
|
+
// 64-bit id survives JS number precision, null to start a fresh session/clone.
|
|
217
|
+
session_id: sessionId ?? null,
|
|
218
|
+
};
|
|
219
|
+
if (idempotencyKey) {
|
|
220
|
+
body.idempotency_key = idempotencyKey;
|
|
221
|
+
}
|
|
222
|
+
return callRpc<JoeSubmitResponse>({
|
|
223
|
+
apiKey,
|
|
224
|
+
apiBaseUrl,
|
|
225
|
+
fn: "joe_command_submit",
|
|
226
|
+
body,
|
|
227
|
+
operation: `Failed to submit ${command} command`,
|
|
228
|
+
debug,
|
|
229
|
+
});
|
|
230
|
+
}
|
|
231
|
+
|
|
232
|
+
export interface CommandIdParams {
|
|
233
|
+
apiKey: string;
|
|
234
|
+
apiBaseUrl: string;
|
|
235
|
+
commandId: string;
|
|
236
|
+
debug?: boolean;
|
|
237
|
+
}
|
|
238
|
+
|
|
239
|
+
/** Poll a command's status metadata (`v1.joe_command_status`). */
|
|
240
|
+
export async function getCommandStatus(params: CommandIdParams): Promise<JoeStatusResponse> {
|
|
241
|
+
const { apiKey, apiBaseUrl, commandId, debug } = params;
|
|
242
|
+
if (!commandId) {
|
|
243
|
+
throw new Error("commandId is required");
|
|
244
|
+
}
|
|
245
|
+
return callRpc<JoeStatusResponse>({
|
|
246
|
+
apiKey,
|
|
247
|
+
apiBaseUrl,
|
|
248
|
+
fn: "joe_command_status",
|
|
249
|
+
body: { command_id: commandId },
|
|
250
|
+
operation: "Failed to fetch command status",
|
|
251
|
+
debug,
|
|
252
|
+
});
|
|
253
|
+
}
|
|
254
|
+
|
|
255
|
+
/** Fetch a command's result body (`v1.joe_command_result`). */
|
|
256
|
+
export async function getCommandResult(params: CommandIdParams): Promise<JoeCommandResult> {
|
|
257
|
+
const { apiKey, apiBaseUrl, commandId, debug } = params;
|
|
258
|
+
if (!commandId) {
|
|
259
|
+
throw new Error("commandId is required");
|
|
260
|
+
}
|
|
261
|
+
return callRpc<JoeCommandResult>({
|
|
262
|
+
apiKey,
|
|
263
|
+
apiBaseUrl,
|
|
264
|
+
fn: "joe_command_result",
|
|
265
|
+
body: { command_id: commandId },
|
|
266
|
+
operation: "Failed to fetch command result",
|
|
267
|
+
debug,
|
|
268
|
+
});
|
|
269
|
+
}
|
|
270
|
+
|
|
271
|
+
export interface ListProjectsParams {
|
|
272
|
+
apiKey: string;
|
|
273
|
+
apiBaseUrl: string;
|
|
274
|
+
orgId?: number;
|
|
275
|
+
debug?: boolean;
|
|
276
|
+
}
|
|
277
|
+
|
|
278
|
+
interface RawProjectRow {
|
|
279
|
+
id?: number;
|
|
280
|
+
project_id?: number;
|
|
281
|
+
alias?: string | null;
|
|
282
|
+
name?: string | null;
|
|
283
|
+
label?: string | null;
|
|
284
|
+
joe_ready?: boolean;
|
|
285
|
+
joe_api_v2_enabled?: boolean;
|
|
286
|
+
tunnel?: boolean;
|
|
287
|
+
tunnel_ready?: boolean;
|
|
288
|
+
}
|
|
289
|
+
|
|
290
|
+
function normalizeProjectRow(row: RawProjectRow): ProjectListItem {
|
|
291
|
+
const projectId = row.project_id ?? row.id ?? 0;
|
|
292
|
+
return {
|
|
293
|
+
project_id: Number(projectId),
|
|
294
|
+
alias: row.alias ?? null,
|
|
295
|
+
name: row.name ?? null,
|
|
296
|
+
label: row.label ?? null,
|
|
297
|
+
joe_ready: Boolean(row.joe_ready ?? row.joe_api_v2_enabled ?? false),
|
|
298
|
+
tunnel: Boolean(row.tunnel ?? row.tunnel_ready ?? false),
|
|
299
|
+
};
|
|
300
|
+
}
|
|
301
|
+
|
|
302
|
+
/**
|
|
303
|
+
* List the org's projects (org-level discovery — NOT a Joe endpoint; SPEC §6).
|
|
304
|
+
* Backed by the general platform projects rpc, extended by item S.1 to surface the
|
|
305
|
+
* per-project `joe_ready` + `tunnel` state so an agent knows which it may address.
|
|
306
|
+
*/
|
|
307
|
+
export async function listProjects(params: ListProjectsParams): Promise<ProjectListItem[]> {
|
|
308
|
+
const { apiKey, apiBaseUrl, orgId, debug } = params;
|
|
309
|
+
const body: Record<string, unknown> = {};
|
|
310
|
+
if (typeof orgId === "number") {
|
|
311
|
+
body.org_id = orgId;
|
|
312
|
+
}
|
|
313
|
+
const rows = await callRpc<RawProjectRow[]>({
|
|
314
|
+
apiKey,
|
|
315
|
+
apiBaseUrl,
|
|
316
|
+
fn: "projects_list",
|
|
317
|
+
body,
|
|
318
|
+
operation: "Failed to list projects",
|
|
319
|
+
debug,
|
|
320
|
+
});
|
|
321
|
+
if (!Array.isArray(rows)) {
|
|
322
|
+
return [];
|
|
323
|
+
}
|
|
324
|
+
return rows.map(normalizeProjectRow);
|
|
325
|
+
}
|
|
326
|
+
|
|
327
|
+
// ---------------------------------------------------------------------------
|
|
328
|
+
// Project id-or-alias resolution (S.0)
|
|
329
|
+
// ---------------------------------------------------------------------------
|
|
330
|
+
|
|
331
|
+
/** A bare numeric `--project` value is a project id; anything else is an alias/name. */
|
|
332
|
+
export function isNumericProjectRef(ref: string): boolean {
|
|
333
|
+
return /^[0-9]+$/.test(ref.trim());
|
|
334
|
+
}
|
|
335
|
+
|
|
336
|
+
export interface ResolveProjectParams {
|
|
337
|
+
apiKey: string;
|
|
338
|
+
apiBaseUrl: string;
|
|
339
|
+
project: string;
|
|
340
|
+
orgId?: number;
|
|
341
|
+
debug?: boolean;
|
|
342
|
+
}
|
|
343
|
+
|
|
344
|
+
/**
|
|
345
|
+
* Resolve `--project <id|alias>` to a numeric project id. A numeric ref is used
|
|
346
|
+
* directly (no lookup); an alias/name is resolved against the projects list, so
|
|
347
|
+
* `--project main-db` ≡ `--project 12` (SPEC §6, finding S.0).
|
|
348
|
+
*/
|
|
349
|
+
export async function resolveProjectId(params: ResolveProjectParams): Promise<number> {
|
|
350
|
+
// TODO(joe-v2): unify --project resolution once projects_list (S.1) returns the dblab instance id
|
|
351
|
+
const ref = String(params.project ?? "").trim();
|
|
352
|
+
if (!ref) {
|
|
353
|
+
throw new Error("project is required (--project <id|alias>)");
|
|
354
|
+
}
|
|
355
|
+
if (isNumericProjectRef(ref)) {
|
|
356
|
+
return Number(ref);
|
|
357
|
+
}
|
|
358
|
+
const projects = await listProjects({
|
|
359
|
+
apiKey: params.apiKey,
|
|
360
|
+
apiBaseUrl: params.apiBaseUrl,
|
|
361
|
+
orgId: params.orgId,
|
|
362
|
+
debug: params.debug,
|
|
363
|
+
});
|
|
364
|
+
const needle = ref.toLowerCase();
|
|
365
|
+
const match = projects.find(
|
|
366
|
+
(p) =>
|
|
367
|
+
(p.alias !== null && p.alias.toLowerCase() === needle) ||
|
|
368
|
+
(p.name !== null && p.name.toLowerCase() === needle)
|
|
369
|
+
);
|
|
370
|
+
if (!match) {
|
|
371
|
+
throw new Error(
|
|
372
|
+
`Project not found for alias/name '${ref}'. Run 'pgai projects' to see available projects.`
|
|
373
|
+
);
|
|
374
|
+
}
|
|
375
|
+
return match.project_id;
|
|
376
|
+
}
|
|
377
|
+
|
|
378
|
+
// ---------------------------------------------------------------------------
|
|
379
|
+
// Per-project session persistence (auto-threaded optimize loop, finding B-1)
|
|
380
|
+
// ---------------------------------------------------------------------------
|
|
381
|
+
|
|
382
|
+
function sessionStorePath(dir?: string): string {
|
|
383
|
+
return path.join(dir ?? config.getConfigDir(), "joe-sessions.json");
|
|
384
|
+
}
|
|
385
|
+
|
|
386
|
+
function readSessionStore(dir?: string): Record<string, string> {
|
|
387
|
+
const file = sessionStorePath(dir);
|
|
388
|
+
if (!fs.existsSync(file)) {
|
|
389
|
+
return {};
|
|
390
|
+
}
|
|
391
|
+
try {
|
|
392
|
+
const parsed = JSON.parse(fs.readFileSync(file, "utf8")) as unknown;
|
|
393
|
+
if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
|
|
394
|
+
return parsed as Record<string, string>;
|
|
395
|
+
}
|
|
396
|
+
} catch {
|
|
397
|
+
// Corrupt store — treat as empty rather than failing a command.
|
|
398
|
+
}
|
|
399
|
+
return {};
|
|
400
|
+
}
|
|
401
|
+
|
|
402
|
+
export function readStoredSessionId(projectId: number, dir?: string): string | null {
|
|
403
|
+
const store = readSessionStore(dir);
|
|
404
|
+
const value = store[String(projectId)];
|
|
405
|
+
return typeof value === "string" && value.length > 0 ? value : null;
|
|
406
|
+
}
|
|
407
|
+
|
|
408
|
+
export function writeStoredSessionId(projectId: number, sessionId: string, dir?: string): void {
|
|
409
|
+
const baseDir = dir ?? config.getConfigDir();
|
|
410
|
+
if (!fs.existsSync(baseDir)) {
|
|
411
|
+
fs.mkdirSync(baseDir, { recursive: true, mode: 0o700 });
|
|
412
|
+
}
|
|
413
|
+
const store = readSessionStore(dir);
|
|
414
|
+
store[String(projectId)] = sessionId;
|
|
415
|
+
fs.writeFileSync(sessionStorePath(dir), JSON.stringify(store, null, 2) + "\n", { mode: 0o600 });
|
|
416
|
+
}
|
|
417
|
+
|
|
418
|
+
export function clearStoredSessionId(projectId: number, dir?: string): void {
|
|
419
|
+
const file = sessionStorePath(dir);
|
|
420
|
+
if (!fs.existsSync(file)) {
|
|
421
|
+
return;
|
|
422
|
+
}
|
|
423
|
+
const store = readSessionStore(dir);
|
|
424
|
+
if (String(projectId) in store) {
|
|
425
|
+
delete store[String(projectId)];
|
|
426
|
+
fs.writeFileSync(file, JSON.stringify(store, null, 2) + "\n", { mode: 0o600 });
|
|
427
|
+
}
|
|
428
|
+
}
|
|
429
|
+
|
|
430
|
+
// ---------------------------------------------------------------------------
|
|
431
|
+
// Idempotency keys
|
|
432
|
+
// ---------------------------------------------------------------------------
|
|
433
|
+
|
|
434
|
+
/**
|
|
435
|
+
* Compute the idempotency key for a submit. Executing commands (`exec`/`explain`)
|
|
436
|
+
* mint a STABLE key derived from `(project_id, command, sql, args)` so a
|
|
437
|
+
* timeout-resubmit dedupes back to the original command_id (at-most-once — SPEC M8);
|
|
438
|
+
* non-executing commands generate a fresh per-invocation key.
|
|
439
|
+
*/
|
|
440
|
+
export function computeIdempotencyKey(
|
|
441
|
+
command: JoeCommand,
|
|
442
|
+
projectId: number,
|
|
443
|
+
sql: string | null | undefined,
|
|
444
|
+
args: Record<string, unknown> | null | undefined
|
|
445
|
+
): string {
|
|
446
|
+
if (EXECUTING_COMMANDS.has(command)) {
|
|
447
|
+
const canonical = `${projectId}\n${command}\n${sql ?? ""}\n${JSON.stringify(args ?? null)}`;
|
|
448
|
+
return `joe-${crypto.createHash("sha256").update(canonical).digest("hex")}`;
|
|
449
|
+
}
|
|
450
|
+
return `joe-${crypto.randomUUID()}`;
|
|
451
|
+
}
|
|
452
|
+
|
|
453
|
+
// ---------------------------------------------------------------------------
|
|
454
|
+
// Submit-then-poll one-shot
|
|
455
|
+
// ---------------------------------------------------------------------------
|
|
456
|
+
|
|
457
|
+
export interface RunCommandParams {
|
|
458
|
+
apiKey: string;
|
|
459
|
+
apiBaseUrl: string;
|
|
460
|
+
command: JoeCommand;
|
|
461
|
+
projectId: number;
|
|
462
|
+
sql?: string | null;
|
|
463
|
+
args?: Record<string, unknown> | null;
|
|
464
|
+
sessionId?: string | null;
|
|
465
|
+
idempotencyKey?: string | null;
|
|
466
|
+
budgetMs?: number;
|
|
467
|
+
pollIntervalMs?: number;
|
|
468
|
+
debug?: boolean;
|
|
469
|
+
/** Injectable clock/sleep for deterministic tests. */
|
|
470
|
+
now?: () => number;
|
|
471
|
+
sleep?: (ms: number) => Promise<void>;
|
|
472
|
+
}
|
|
473
|
+
|
|
474
|
+
export interface RunCommandOutcome {
|
|
475
|
+
commandId: string;
|
|
476
|
+
/** The session id returned by submit (thread back / persist for the next command). */
|
|
477
|
+
sessionId: string | null;
|
|
478
|
+
status: JoeStatus;
|
|
479
|
+
/** Populated once the command reaches a terminal state (done/error/timed_out). */
|
|
480
|
+
result: JoeCommandResult | null;
|
|
481
|
+
/** True when the ≤ budget one-shot expired before a terminal state — resume by id. */
|
|
482
|
+
budgetExpired: boolean;
|
|
483
|
+
}
|
|
484
|
+
|
|
485
|
+
const defaultSleep = (ms: number): Promise<void> => new Promise((r) => setTimeout(r, ms));
|
|
486
|
+
|
|
487
|
+
/**
|
|
488
|
+
* Submit a command then client-side poll status/result within the one-shot budget.
|
|
489
|
+
* On a terminal state returns the result; on budget expiry returns a resume handle
|
|
490
|
+
* (`budgetExpired: true`) so the caller can `pgai result <command_id>` later.
|
|
491
|
+
*/
|
|
492
|
+
export async function runCommand(params: RunCommandParams): Promise<RunCommandOutcome> {
|
|
493
|
+
const {
|
|
494
|
+
apiKey,
|
|
495
|
+
apiBaseUrl,
|
|
496
|
+
command,
|
|
497
|
+
projectId,
|
|
498
|
+
sql,
|
|
499
|
+
args,
|
|
500
|
+
sessionId,
|
|
501
|
+
debug,
|
|
502
|
+
} = params;
|
|
503
|
+
const budgetMs = params.budgetMs ?? DEFAULT_BUDGET_MS;
|
|
504
|
+
const pollIntervalMs = params.pollIntervalMs ?? DEFAULT_POLL_INTERVAL_MS;
|
|
505
|
+
const now = params.now ?? Date.now;
|
|
506
|
+
const sleep = params.sleep ?? defaultSleep;
|
|
507
|
+
const idempotencyKey =
|
|
508
|
+
params.idempotencyKey ?? computeIdempotencyKey(command, projectId, sql, args);
|
|
509
|
+
|
|
510
|
+
const submitted = await submitCommand({
|
|
511
|
+
apiKey,
|
|
512
|
+
apiBaseUrl,
|
|
513
|
+
command,
|
|
514
|
+
projectId,
|
|
515
|
+
sql,
|
|
516
|
+
args,
|
|
517
|
+
sessionId,
|
|
518
|
+
idempotencyKey,
|
|
519
|
+
debug,
|
|
520
|
+
});
|
|
521
|
+
|
|
522
|
+
const commandId = submitted.command_id;
|
|
523
|
+
const newSessionId = submitted.session_id ?? sessionId ?? null;
|
|
524
|
+
const deadline = now() + budgetMs;
|
|
525
|
+
let status: JoeStatus = submitted.status;
|
|
526
|
+
|
|
527
|
+
// Poll status until terminal or the one-shot budget is exhausted.
|
|
528
|
+
// eslint-disable-next-line no-constant-condition
|
|
529
|
+
while (true) {
|
|
530
|
+
const statusResp = await getCommandStatus({ apiKey, apiBaseUrl, commandId, debug });
|
|
531
|
+
status = statusResp.status;
|
|
532
|
+
if (TERMINAL_STATES.has(status)) {
|
|
533
|
+
break;
|
|
534
|
+
}
|
|
535
|
+
if (now() >= deadline) {
|
|
536
|
+
return { commandId, sessionId: newSessionId, status, result: null, budgetExpired: true };
|
|
537
|
+
}
|
|
538
|
+
await sleep(pollIntervalMs);
|
|
539
|
+
}
|
|
540
|
+
|
|
541
|
+
const result = await getCommandResult({ apiKey, apiBaseUrl, commandId, debug });
|
|
542
|
+
return { commandId, sessionId: newSessionId, status, result, budgetExpired: false };
|
|
543
|
+
}
|
|
544
|
+
|
|
545
|
+
// ---------------------------------------------------------------------------
|
|
546
|
+
// High-level orchestrator shared by the CLI and MCP surfaces
|
|
547
|
+
// ---------------------------------------------------------------------------
|
|
548
|
+
|
|
549
|
+
export interface ExecuteJoeParams {
|
|
550
|
+
apiKey: string;
|
|
551
|
+
apiBaseUrl: string;
|
|
552
|
+
command: JoeCommand;
|
|
553
|
+
/** Raw `--project <id|alias>` value. */
|
|
554
|
+
project: string;
|
|
555
|
+
sql?: string | null;
|
|
556
|
+
args?: Record<string, unknown> | null;
|
|
557
|
+
/** Explicit `--session <id>` (overrides the stored/auto session). */
|
|
558
|
+
session?: string | null;
|
|
559
|
+
/** `--new-session` — force a fresh clone (null session). */
|
|
560
|
+
newSession?: boolean;
|
|
561
|
+
orgId?: number;
|
|
562
|
+
budgetMs?: number;
|
|
563
|
+
pollIntervalMs?: number;
|
|
564
|
+
debug?: boolean;
|
|
565
|
+
/** Session-store directory (defaults to the user config dir). */
|
|
566
|
+
sessionDir?: string;
|
|
567
|
+
now?: () => number;
|
|
568
|
+
sleep?: (ms: number) => Promise<void>;
|
|
569
|
+
}
|
|
570
|
+
|
|
571
|
+
export interface ExecuteJoeOutcome extends RunCommandOutcome {
|
|
572
|
+
command: JoeCommand;
|
|
573
|
+
projectId: number;
|
|
574
|
+
}
|
|
575
|
+
|
|
576
|
+
/**
|
|
577
|
+
* Resolve the project (id-or-alias), pick the session (explicit / stored / new),
|
|
578
|
+
* run the submit-then-poll one-shot, and persist the returned session for the next
|
|
579
|
+
* command on the same project — the auto-threaded optimize loop.
|
|
580
|
+
*/
|
|
581
|
+
export async function executeJoeCommand(params: ExecuteJoeParams): Promise<ExecuteJoeOutcome> {
|
|
582
|
+
const projectId = await resolveProjectId({
|
|
583
|
+
apiKey: params.apiKey,
|
|
584
|
+
apiBaseUrl: params.apiBaseUrl,
|
|
585
|
+
project: params.project,
|
|
586
|
+
orgId: params.orgId,
|
|
587
|
+
debug: params.debug,
|
|
588
|
+
});
|
|
589
|
+
|
|
590
|
+
let sessionId: string | null;
|
|
591
|
+
if (params.newSession) {
|
|
592
|
+
clearStoredSessionId(projectId, params.sessionDir);
|
|
593
|
+
sessionId = null;
|
|
594
|
+
} else if (params.session) {
|
|
595
|
+
sessionId = String(params.session);
|
|
596
|
+
} else {
|
|
597
|
+
sessionId = readStoredSessionId(projectId, params.sessionDir);
|
|
598
|
+
}
|
|
599
|
+
|
|
600
|
+
const outcome = await runCommand({
|
|
601
|
+
apiKey: params.apiKey,
|
|
602
|
+
apiBaseUrl: params.apiBaseUrl,
|
|
603
|
+
command: params.command,
|
|
604
|
+
projectId,
|
|
605
|
+
sql: params.sql,
|
|
606
|
+
args: params.args,
|
|
607
|
+
sessionId,
|
|
608
|
+
budgetMs: params.budgetMs,
|
|
609
|
+
pollIntervalMs: params.pollIntervalMs,
|
|
610
|
+
debug: params.debug,
|
|
611
|
+
now: params.now,
|
|
612
|
+
sleep: params.sleep,
|
|
613
|
+
});
|
|
614
|
+
|
|
615
|
+
if (outcome.sessionId) {
|
|
616
|
+
writeStoredSessionId(projectId, outcome.sessionId, params.sessionDir);
|
|
617
|
+
}
|
|
618
|
+
|
|
619
|
+
return { ...outcome, command: params.command, projectId };
|
|
620
|
+
}
|
|
621
|
+
|
|
622
|
+
// ---------------------------------------------------------------------------
|
|
623
|
+
// Presentation helpers (pure — unit tested)
|
|
624
|
+
// ---------------------------------------------------------------------------
|
|
625
|
+
|
|
626
|
+
interface PlanNode {
|
|
627
|
+
"Node Type"?: string;
|
|
628
|
+
"Relation Name"?: string;
|
|
629
|
+
Plans?: PlanNode[];
|
|
630
|
+
[key: string]: unknown;
|
|
631
|
+
}
|
|
632
|
+
|
|
633
|
+
/**
|
|
634
|
+
* Lightweight CLIENT-SIDE plan flagging (SPEC §6). The API returns no
|
|
635
|
+
* server-authored recommendations (dropped as an injection vector), so the CLI
|
|
636
|
+
* flags obvious issues (e.g. a Seq Scan) from the returned plan_json itself.
|
|
637
|
+
*/
|
|
638
|
+
export function clientSidePlanFlags(planJson: unknown): string[] {
|
|
639
|
+
const flags: string[] = [];
|
|
640
|
+
const walk = (node: PlanNode | undefined): void => {
|
|
641
|
+
if (!node || typeof node !== "object") {
|
|
642
|
+
return;
|
|
643
|
+
}
|
|
644
|
+
const nodeType = node["Node Type"];
|
|
645
|
+
if (nodeType === "Seq Scan") {
|
|
646
|
+
const rel = node["Relation Name"];
|
|
647
|
+
flags.push(
|
|
648
|
+
`client-side: Seq Scan${rel ? ` on ${rel}` : ""} — no index serves this predicate; consider adding one.`
|
|
649
|
+
);
|
|
650
|
+
}
|
|
651
|
+
if (Array.isArray(node.Plans)) {
|
|
652
|
+
for (const child of node.Plans) {
|
|
653
|
+
walk(child);
|
|
654
|
+
}
|
|
655
|
+
}
|
|
656
|
+
};
|
|
657
|
+
if (planJson && typeof planJson === "object") {
|
|
658
|
+
const root = planJson as { Plan?: PlanNode };
|
|
659
|
+
walk(root.Plan ?? (planJson as PlanNode));
|
|
660
|
+
}
|
|
661
|
+
return flags;
|
|
662
|
+
}
|
|
663
|
+
|
|
664
|
+
/** Format a terminal command result as human-readable text (non-JSON mode). */
|
|
665
|
+
export function formatJoeResult(command: JoeCommand, result: JoeCommandResult): string {
|
|
666
|
+
const lines: string[] = [];
|
|
667
|
+
switch (command) {
|
|
668
|
+
case "plan":
|
|
669
|
+
case "explain": {
|
|
670
|
+
if (result.plan_text) {
|
|
671
|
+
lines.push(result.plan_text);
|
|
672
|
+
}
|
|
673
|
+
for (const flag of clientSidePlanFlags(result.plan_json)) {
|
|
674
|
+
lines.push(`⚑ ${flag}`);
|
|
675
|
+
}
|
|
676
|
+
if (result.queryid || result.plan_fingerprint) {
|
|
677
|
+
lines.push(`(queryid ${result.queryid ?? "?"} · plan_fp ${result.plan_fingerprint ?? "?"})`);
|
|
678
|
+
}
|
|
679
|
+
break;
|
|
680
|
+
}
|
|
681
|
+
case "exec": {
|
|
682
|
+
const notices = result.notices ?? [];
|
|
683
|
+
const rowCount = result.row_count ?? 0;
|
|
684
|
+
lines.push(`${notices.join(" ") || "OK"} · ${rowCount} row${rowCount === 1 ? "" : "s"}`);
|
|
685
|
+
if (result.result_rows && result.result_rows.length > 0) {
|
|
686
|
+
lines.push(JSON.stringify(result.result_rows, null, 2));
|
|
687
|
+
}
|
|
688
|
+
break;
|
|
689
|
+
}
|
|
690
|
+
case "hypo": {
|
|
691
|
+
lines.push(result.hypo_used ? "hypothetical index would be used" : "hypothetical index would NOT be used");
|
|
692
|
+
if (result.hypo_plan !== undefined) {
|
|
693
|
+
lines.push(JSON.stringify(result.hypo_plan, null, 2));
|
|
694
|
+
}
|
|
695
|
+
break;
|
|
696
|
+
}
|
|
697
|
+
case "activity":
|
|
698
|
+
case "describe": {
|
|
699
|
+
lines.push(JSON.stringify(result.snapshot ?? null, null, 2));
|
|
700
|
+
break;
|
|
701
|
+
}
|
|
702
|
+
case "terminate": {
|
|
703
|
+
lines.push(`terminated ${result.terminated ? "yes" : "no"} · pid ${result.pid ?? "?"}`);
|
|
704
|
+
break;
|
|
705
|
+
}
|
|
706
|
+
case "reset": {
|
|
707
|
+
lines.push(`reset ${result.reset ? "ok" : "no"}`);
|
|
708
|
+
break;
|
|
709
|
+
}
|
|
710
|
+
}
|
|
711
|
+
return lines.join("\n");
|
|
712
|
+
}
|
|
713
|
+
|
|
714
|
+
/** Render `pgai projects` as the brief's fixed-width table. */
|
|
715
|
+
export function formatProjectsTable(projects: ProjectListItem[]): string {
|
|
716
|
+
const header = ["PROJECT_ID", "ALIAS", "PROJECT", "JOE", "TUNNEL"];
|
|
717
|
+
const rows = projects.map((p) => [
|
|
718
|
+
String(p.project_id),
|
|
719
|
+
p.alias ?? "-",
|
|
720
|
+
p.name ?? "-",
|
|
721
|
+
p.joe_ready ? "ready" : "no",
|
|
722
|
+
p.tunnel ? "yes" : "no",
|
|
723
|
+
]);
|
|
724
|
+
const widths = header.map((h, i) =>
|
|
725
|
+
Math.max(h.length, ...rows.map((r) => r[i].length), 0)
|
|
726
|
+
);
|
|
727
|
+
const pad = (cells: string[]): string =>
|
|
728
|
+
cells.map((c, i) => c.padEnd(i === cells.length - 1 ? 0 : widths[i])).join(" ").trimEnd();
|
|
729
|
+
return [pad(header), ...rows.map(pad)].join("\n");
|
|
730
|
+
}
|