@clawdbuilder/jig-bedrock 0.0.0-stage → 0.1.1
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/LICENSE +9 -0
- package/README.md +172 -2
- package/bin/windows-browser +4 -0
- package/index.ts +3 -0
- package/package.json +51 -5
- package/src/aws-discovery.ts +226 -0
- package/src/bedrock-fetch.ts +77 -0
- package/src/budget.ts +129 -0
- package/src/cache.ts +22 -0
- package/src/catalog.ts +514 -0
- package/src/check.ts +14 -0
- package/src/child.ts +7 -0
- package/src/commands.ts +344 -0
- package/src/config.ts +214 -0
- package/src/doctor.ts +172 -0
- package/src/events.ts +46 -0
- package/src/identity.ts +3 -0
- package/src/index.ts +33 -0
- package/src/jig/client.ts +63 -0
- package/src/jig/types.ts +12 -0
- package/src/login.ts +176 -0
- package/src/messages.ts +45 -0
- package/src/pi-host.ts +54 -0
- package/src/prompt-cache.ts +50 -0
- package/src/provider.ts +344 -0
- package/src/sigv4.ts +92 -0
- package/src/sso-check.ts +111 -0
- package/src/sso.ts +339 -0
- package/src/start-sign-in.ts +53 -0
- package/src/statistics.ts +315 -0
- package/src/switch-cost.ts +91 -0
- package/src/usage.ts +138 -0
- package/src/wsl.ts +107 -0
package/src/budget.ts
ADDED
|
@@ -0,0 +1,129 @@
|
|
|
1
|
+
import { homedir } from "node:os";
|
|
2
|
+
import { dirname, join } from "node:path";
|
|
3
|
+
import { atomicCache } from "./cache.ts";
|
|
4
|
+
import { boundedRegular } from "./catalog.ts";
|
|
5
|
+
import type { ResolvedConfig } from "./config.ts";
|
|
6
|
+
import { createStatistics, parseUsage, pricedSpend, type StatisticsSource } from "./statistics.ts";
|
|
7
|
+
import { runUsage, type UsageRunner } from "./usage.ts";
|
|
8
|
+
/** spentUsd sums the priced rows; unpriced says some usage had no price, so the real spend is higher. */
|
|
9
|
+
export interface Budget {
|
|
10
|
+
monthUtc: string;
|
|
11
|
+
spentUsd: number;
|
|
12
|
+
unpriced: boolean;
|
|
13
|
+
limitUsd?: number;
|
|
14
|
+
remainingUsd?: number;
|
|
15
|
+
observedAt: number;
|
|
16
|
+
refreshFailed: boolean;
|
|
17
|
+
}
|
|
18
|
+
export interface BudgetSource {
|
|
19
|
+
load(): Budget | undefined;
|
|
20
|
+
refresh(options?: { force?: boolean; signal?: AbortSignal }): Promise<void>;
|
|
21
|
+
}
|
|
22
|
+
export function parseBudget(output: string, monthUtc: string, observedAt: number, limitUsd?: number): Budget {
|
|
23
|
+
if (new Date(observedAt).toISOString().slice(0, 7) !== monthUtc)
|
|
24
|
+
throw new Error("usage program monthly result unavailable");
|
|
25
|
+
const { spentUsd, unpriced } = pricedSpend(parseUsage(output, "monthly"));
|
|
26
|
+
return {
|
|
27
|
+
monthUtc,
|
|
28
|
+
spentUsd,
|
|
29
|
+
unpriced,
|
|
30
|
+
observedAt,
|
|
31
|
+
refreshFailed: false,
|
|
32
|
+
...(limitUsd === undefined ? {} : { limitUsd, remainingUsd: limitUsd - spentUsd }),
|
|
33
|
+
};
|
|
34
|
+
}
|
|
35
|
+
export function createBudget(
|
|
36
|
+
config: ResolvedConfig,
|
|
37
|
+
runner: UsageRunner = runUsage,
|
|
38
|
+
file = join(homedir(), ".cache/jig-bedrock/budget.json"),
|
|
39
|
+
now = Date.now,
|
|
40
|
+
statistics: StatisticsSource = createStatistics(config, runner, dirname(file), now),
|
|
41
|
+
): BudgetSource {
|
|
42
|
+
const scope = JSON.stringify({
|
|
43
|
+
profile: config.profile,
|
|
44
|
+
team: config.team,
|
|
45
|
+
regions: [...config.regions].sort(),
|
|
46
|
+
usageDir: config.usageDir,
|
|
47
|
+
});
|
|
48
|
+
let flight: Promise<void> | undefined,
|
|
49
|
+
lastAttempt = -Infinity;
|
|
50
|
+
const raw = () => {
|
|
51
|
+
try {
|
|
52
|
+
return JSON.parse(boundedRegular(file, 32768, true) || "{}");
|
|
53
|
+
} catch {
|
|
54
|
+
return {};
|
|
55
|
+
}
|
|
56
|
+
};
|
|
57
|
+
const load = (): Budget | undefined => {
|
|
58
|
+
const c = raw();
|
|
59
|
+
if (
|
|
60
|
+
c.schema !== 1 ||
|
|
61
|
+
c.profile !== config.profile ||
|
|
62
|
+
c.scope !== scope ||
|
|
63
|
+
c.monthUtc !== new Date(now()).toISOString().slice(0, 7) ||
|
|
64
|
+
!Number.isFinite(c.spentUsd) ||
|
|
65
|
+
c.spentUsd < 0 ||
|
|
66
|
+
!Number.isSafeInteger(c.observedAt) ||
|
|
67
|
+
c.observedAt < 0 ||
|
|
68
|
+
c.observedAt > now()
|
|
69
|
+
)
|
|
70
|
+
return;
|
|
71
|
+
return {
|
|
72
|
+
monthUtc: c.monthUtc,
|
|
73
|
+
spentUsd: c.spentUsd,
|
|
74
|
+
unpriced: c.unpriced === true,
|
|
75
|
+
observedAt: c.observedAt,
|
|
76
|
+
refreshFailed: c.refreshFailed === true,
|
|
77
|
+
...(config.monthlyBudgetUsd === undefined
|
|
78
|
+
? {}
|
|
79
|
+
: { limitUsd: config.monthlyBudgetUsd, remainingUsd: config.monthlyBudgetUsd - c.spentUsd }),
|
|
80
|
+
};
|
|
81
|
+
};
|
|
82
|
+
const fromSnapshot = (snapshot: ReturnType<StatisticsSource["load"]>) =>
|
|
83
|
+
snapshot && { ...pricedSpend(snapshot.report), observedAt: snapshot.observedAt };
|
|
84
|
+
const perform = async (options: { force?: boolean; signal?: AbortSignal } = {}) => {
|
|
85
|
+
const old = raw(),
|
|
86
|
+
started = now(),
|
|
87
|
+
monthUtc = new Date(started).toISOString().slice(0, 7);
|
|
88
|
+
if (
|
|
89
|
+
!options.force &&
|
|
90
|
+
started - Math.max(lastAttempt, old.scope === scope ? old.lastAttemptAt || 0 : 0) < 300_000
|
|
91
|
+
)
|
|
92
|
+
return;
|
|
93
|
+
lastAttempt = started;
|
|
94
|
+
try {
|
|
95
|
+
const snapshot = await statistics.refresh("monthly", { signal: options.signal });
|
|
96
|
+
atomicCache(file, {
|
|
97
|
+
schema: 1,
|
|
98
|
+
scope,
|
|
99
|
+
profile: config.profile,
|
|
100
|
+
monthUtc,
|
|
101
|
+
lastAttemptAt: started,
|
|
102
|
+
observedAt: snapshot.observedAt,
|
|
103
|
+
refreshFailed: false,
|
|
104
|
+
...pricedSpend(snapshot.report),
|
|
105
|
+
});
|
|
106
|
+
} catch {
|
|
107
|
+
// Keep the last known figure: this file's own, or else the last valid usage program report the statistics cache still holds.
|
|
108
|
+
const known = load() ?? fromSnapshot(statistics.load("monthly"));
|
|
109
|
+
try {
|
|
110
|
+
atomicCache(file, {
|
|
111
|
+
schema: 1,
|
|
112
|
+
scope,
|
|
113
|
+
profile: config.profile,
|
|
114
|
+
monthUtc,
|
|
115
|
+
lastAttemptAt: started,
|
|
116
|
+
refreshFailed: true,
|
|
117
|
+
...(known && { spentUsd: known.spentUsd, unpriced: known.unpriced, observedAt: known.observedAt }),
|
|
118
|
+
});
|
|
119
|
+
} catch {}
|
|
120
|
+
}
|
|
121
|
+
};
|
|
122
|
+
return {
|
|
123
|
+
load,
|
|
124
|
+
refresh: (options) =>
|
|
125
|
+
(flight ??= perform(options).finally(() => {
|
|
126
|
+
flight = undefined;
|
|
127
|
+
})),
|
|
128
|
+
};
|
|
129
|
+
}
|
package/src/cache.ts
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
import { randomUUID } from "node:crypto";
|
|
2
|
+
import { lstatSync, mkdirSync, renameSync, rmSync, writeFileSync } from "node:fs";
|
|
3
|
+
import { dirname } from "node:path";
|
|
4
|
+
export function atomicCache(file: string, value: unknown): void {
|
|
5
|
+
const data = JSON.stringify(value);
|
|
6
|
+
if (Buffer.byteLength(data) > 1_048_576) throw new Error("Cache too large");
|
|
7
|
+
for (let p = dirname(file); dirname(p) !== p; p = dirname(p)) {
|
|
8
|
+
try {
|
|
9
|
+
if (lstatSync(p).isSymbolicLink()) throw new Error("Unsafe cache directory");
|
|
10
|
+
} catch (e) {
|
|
11
|
+
if ((e as NodeJS.ErrnoException).code !== "ENOENT") throw e;
|
|
12
|
+
}
|
|
13
|
+
}
|
|
14
|
+
mkdirSync(dirname(file), { recursive: true, mode: 0o700 });
|
|
15
|
+
const temp = `${file}.${randomUUID()}.tmp`;
|
|
16
|
+
try {
|
|
17
|
+
writeFileSync(temp, data, { flag: "wx", mode: 0o600 });
|
|
18
|
+
renameSync(temp, file);
|
|
19
|
+
} finally {
|
|
20
|
+
rmSync(temp, { force: true });
|
|
21
|
+
}
|
|
22
|
+
}
|
package/src/catalog.ts
ADDED
|
@@ -0,0 +1,514 @@
|
|
|
1
|
+
import { closeSync, constants, fstatSync, lstatSync, openSync, readFileSync, readSync } from "node:fs";
|
|
2
|
+
import { homedir } from "node:os";
|
|
3
|
+
import { dirname, join, resolve } from "node:path";
|
|
4
|
+
import type {
|
|
5
|
+
ExtensionAPI,
|
|
6
|
+
ExtensionContext,
|
|
7
|
+
ProviderConfig,
|
|
8
|
+
ProviderModelConfig,
|
|
9
|
+
} from "@earendil-works/pi-coding-agent";
|
|
10
|
+
import { isJigBedrock, JIG_BEDROCK } from "./identity.ts";
|
|
11
|
+
import { isAnthropicModel, noThinkingOff, rewriteThinkingOff, usesAdaptiveThinking } from "./messages.ts";
|
|
12
|
+
import { promptCacheFor } from "./prompt-cache.ts";
|
|
13
|
+
|
|
14
|
+
// biome-ignore lint/suspicious/noExplicitAny: Pi's model registry and session entries, checked field by field where they are used
|
|
15
|
+
type Json = any;
|
|
16
|
+
|
|
17
|
+
export { isJigBedrock, JIG_BEDROCK } from "./identity.ts";
|
|
18
|
+
|
|
19
|
+
let activeRegistry: Record<string, Json> = {};
|
|
20
|
+
export function usePiRegistry(native: import("@earendil-works/pi-ai").Provider): void {
|
|
21
|
+
const models = native.getModels();
|
|
22
|
+
activeRegistry = Object.fromEntries(
|
|
23
|
+
models
|
|
24
|
+
.filter((m: Json) => m.provider === "amazon-bedrock" && typeof m.id === "string")
|
|
25
|
+
.map((m: Json) => [m.id, m]),
|
|
26
|
+
);
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
// Pi 1.2 widened ProviderModelConfig to chat, image and classifier models; team profiles are chat models.
|
|
30
|
+
type ChatModelConfig = Extract<ProviderModelConfig, { type?: "chat" }>;
|
|
31
|
+
export type TeamBedrockModel = ChatModelConfig & {
|
|
32
|
+
api: "anthropic-messages" | "bedrock-converse-stream";
|
|
33
|
+
baseUrl: string;
|
|
34
|
+
arn: string;
|
|
35
|
+
foundationModelId: string;
|
|
36
|
+
rateSource: "usage program" | "Pi registry";
|
|
37
|
+
aliases: string[];
|
|
38
|
+
};
|
|
39
|
+
|
|
40
|
+
/**
|
|
41
|
+
* What Pi's Anthropic Messages client is told about Bedrock's route, from Pi's own `getAnthropicCompat` names. The probe on the
|
|
42
|
+
* team account showed native tool changes (the inline-tools beta and the deferred placeholder) working, so those stay on; long
|
|
43
|
+
* cache retention and strict tool schemas are off because the route has not shown them, and the rest is the Anthropic API's
|
|
44
|
+
* default, which the route mirrors. See the report for what was not tested.
|
|
45
|
+
*/
|
|
46
|
+
const MESSAGES_COMPAT = {
|
|
47
|
+
supportsEagerToolInputStreaming: true,
|
|
48
|
+
supportsLongCacheRetention: false,
|
|
49
|
+
supportsCacheControlOnTools: true,
|
|
50
|
+
supportsStrictTools: false,
|
|
51
|
+
supportsTemperature: true,
|
|
52
|
+
sendSessionAffinityHeaders: false,
|
|
53
|
+
allowEmptySignature: false,
|
|
54
|
+
supportsMidConvoSystemMessages: true,
|
|
55
|
+
supportsMidConvoToolChanges: true,
|
|
56
|
+
} as const;
|
|
57
|
+
export type Profile = { arn: string; profile_name: string; foundation_model_id: string; region: string };
|
|
58
|
+
export type Catalog = {
|
|
59
|
+
schema: 1;
|
|
60
|
+
team: string;
|
|
61
|
+
aws_profile: string;
|
|
62
|
+
account_id: string;
|
|
63
|
+
regions: string[];
|
|
64
|
+
fetched_at: number;
|
|
65
|
+
profiles: Profile[];
|
|
66
|
+
};
|
|
67
|
+
const slug = (s: string) =>
|
|
68
|
+
s
|
|
69
|
+
.toLowerCase()
|
|
70
|
+
.replace(/[^a-z0-9]+/g, "-")
|
|
71
|
+
.replace(/^-|-$/g, "");
|
|
72
|
+
const registryEntry = (id: string) =>
|
|
73
|
+
activeRegistry[id] || activeRegistry[`us.${id}`] || activeRegistry[`global.${id}`];
|
|
74
|
+
const registryName = (name: string) => name.replace(/\s+\((?:US|Global)\)$/i, "");
|
|
75
|
+
const label = (s: unknown, n: number): s is string =>
|
|
76
|
+
// biome-ignore lint/suspicious/noControlCharactersInRegex: catalog labels never contain control characters
|
|
77
|
+
typeof s === "string" && s.length > 0 && s.length <= n && !/[\x00-\x1f\x7f]/.test(s);
|
|
78
|
+
export function boundedRegular(file: string, max: number, privateMode = false): string | undefined {
|
|
79
|
+
const fd = openSync(file, constants.O_RDONLY | constants.O_NOFOLLOW | constants.O_NONBLOCK);
|
|
80
|
+
try {
|
|
81
|
+
const meta = fstatSync(fd);
|
|
82
|
+
if (!meta.isFile() || meta.size > max || (privateMode && (meta.mode & 0o777) !== 0o600)) return;
|
|
83
|
+
const buffer = Buffer.alloc(meta.size + 1);
|
|
84
|
+
let total = 0;
|
|
85
|
+
while (total < buffer.length) {
|
|
86
|
+
const count = readSync(fd, buffer, total, buffer.length - total, null);
|
|
87
|
+
if (!count) break;
|
|
88
|
+
total += count;
|
|
89
|
+
}
|
|
90
|
+
return total <= max ? buffer.subarray(0, total).toString("utf8") : undefined;
|
|
91
|
+
} finally {
|
|
92
|
+
closeSync(fd);
|
|
93
|
+
}
|
|
94
|
+
}
|
|
95
|
+
export function validArn(arn: string, region: string, account: string): boolean {
|
|
96
|
+
const p = arn.split(":");
|
|
97
|
+
return (
|
|
98
|
+
p.length === 6 &&
|
|
99
|
+
p[0] === "arn" &&
|
|
100
|
+
["aws", "aws-cn", "aws-us-gov"].includes(p[1] ?? "") &&
|
|
101
|
+
p[2] === "bedrock" &&
|
|
102
|
+
p[3] === region &&
|
|
103
|
+
p[4] === account &&
|
|
104
|
+
/^application-inference-profile\/[A-Za-z0-9-]{1,128}$/.test(p[5] ?? "")
|
|
105
|
+
);
|
|
106
|
+
}
|
|
107
|
+
export function readCatalog(
|
|
108
|
+
file: string,
|
|
109
|
+
expected: { team: string; profile: string; regions: string[] },
|
|
110
|
+
now = Date.now(),
|
|
111
|
+
): Catalog | undefined {
|
|
112
|
+
expected = { ...expected, regions: [...expected.regions].sort() };
|
|
113
|
+
const directory = dirname(file);
|
|
114
|
+
if (
|
|
115
|
+
!label(expected.team, 256) ||
|
|
116
|
+
!label(expected.profile, 128) ||
|
|
117
|
+
!expected.regions.length ||
|
|
118
|
+
new Set(expected.regions).size !== expected.regions.length
|
|
119
|
+
)
|
|
120
|
+
return;
|
|
121
|
+
try {
|
|
122
|
+
for (let current = directory; dirname(current) !== current; current = dirname(current))
|
|
123
|
+
if (lstatSync(current).isSymbolicLink()) return;
|
|
124
|
+
const raw = boundedRegular(file, 1_048_576, true);
|
|
125
|
+
if (!raw) return;
|
|
126
|
+
const c = JSON.parse(raw) as Catalog;
|
|
127
|
+
if (
|
|
128
|
+
c.schema !== 1 ||
|
|
129
|
+
c.team !== expected.team ||
|
|
130
|
+
c.aws_profile !== expected.profile ||
|
|
131
|
+
typeof c.account_id !== "string" ||
|
|
132
|
+
!/^\d{12}$/.test(c.account_id) ||
|
|
133
|
+
!Array.isArray(c.regions) ||
|
|
134
|
+
c.regions.length !== expected.regions.length ||
|
|
135
|
+
c.regions
|
|
136
|
+
.slice()
|
|
137
|
+
.sort()
|
|
138
|
+
.some((r, i) => r !== expected.regions[i]) ||
|
|
139
|
+
!Number.isSafeInteger(c.fetched_at) ||
|
|
140
|
+
c.fetched_at < 0 ||
|
|
141
|
+
c.fetched_at > now / 1000 ||
|
|
142
|
+
!Array.isArray(c.profiles) ||
|
|
143
|
+
c.profiles.length > 256
|
|
144
|
+
)
|
|
145
|
+
return;
|
|
146
|
+
const seen = new Set<string>();
|
|
147
|
+
for (const p of c.profiles) {
|
|
148
|
+
if (
|
|
149
|
+
!p ||
|
|
150
|
+
!expected.regions.includes(p.region) ||
|
|
151
|
+
!label(p.profile_name, 256) ||
|
|
152
|
+
!label(p.foundation_model_id, 256) ||
|
|
153
|
+
typeof p.arn !== "string" ||
|
|
154
|
+
!validArn(p.arn, p.region, c.account_id) ||
|
|
155
|
+
seen.has(p.arn)
|
|
156
|
+
)
|
|
157
|
+
return;
|
|
158
|
+
seen.add(p.arn);
|
|
159
|
+
}
|
|
160
|
+
return c;
|
|
161
|
+
} catch {
|
|
162
|
+
return;
|
|
163
|
+
}
|
|
164
|
+
}
|
|
165
|
+
function usageRates(dir: string): Map<string, ProviderModelConfig["cost"]> {
|
|
166
|
+
try {
|
|
167
|
+
const file = join(dir.startsWith("~/") ? join(homedir(), dir.slice(2)) : resolve(dir), "models.json");
|
|
168
|
+
const raw = boundedRegular(file, 1_048_576);
|
|
169
|
+
if (!raw) return new Map();
|
|
170
|
+
const rows = JSON.parse(raw);
|
|
171
|
+
if (!Array.isArray(rows)) return new Map();
|
|
172
|
+
const rates = new Map<string, ProviderModelConfig["cost"]>();
|
|
173
|
+
for (const row of rows) {
|
|
174
|
+
const p = row?.pricing;
|
|
175
|
+
if (
|
|
176
|
+
label(row?.key, 128) &&
|
|
177
|
+
p &&
|
|
178
|
+
["input", "output", "cache_read", "cache_write"].every((k) => Number.isFinite(p[k]) && p[k] >= 0)
|
|
179
|
+
)
|
|
180
|
+
rates.set(row.key, {
|
|
181
|
+
input: p.input,
|
|
182
|
+
output: p.output,
|
|
183
|
+
cacheRead: p.cache_read,
|
|
184
|
+
cacheWrite: p.cache_write,
|
|
185
|
+
});
|
|
186
|
+
}
|
|
187
|
+
return rates;
|
|
188
|
+
} catch {
|
|
189
|
+
return new Map();
|
|
190
|
+
}
|
|
191
|
+
}
|
|
192
|
+
/** Pi's level map for a model, plus a null off when the model cannot turn thinking off; a null entry hides that level. */
|
|
193
|
+
function thinkingLevelMap(
|
|
194
|
+
base: ChatModelConfig["thinkingLevelMap"],
|
|
195
|
+
foundationModelId: string,
|
|
196
|
+
): Pick<ChatModelConfig, "thinkingLevelMap"> | undefined {
|
|
197
|
+
if (!noThinkingOff(foundationModelId)) return base === undefined ? undefined : { thinkingLevelMap: base };
|
|
198
|
+
return { thinkingLevelMap: { ...(base ?? {}), off: null } };
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
export function discoveredProfiles(
|
|
202
|
+
catalog: Catalog | undefined,
|
|
203
|
+
usageDir: string,
|
|
204
|
+
api: "messages" | "converse" = "messages",
|
|
205
|
+
): { available: TeamBedrockModel[]; unavailable: { profile: string; region: string; reason: string }[] } {
|
|
206
|
+
if (!catalog) return { available: [], unavailable: [] };
|
|
207
|
+
const rates = usageRates(usageDir),
|
|
208
|
+
available: TeamBedrockModel[] = [],
|
|
209
|
+
unavailable = [];
|
|
210
|
+
const grouped = new Map<string, Profile[]>();
|
|
211
|
+
for (const p of catalog.profiles) {
|
|
212
|
+
const r = registryEntry(p.foundation_model_id);
|
|
213
|
+
if (
|
|
214
|
+
!r ||
|
|
215
|
+
!Number.isSafeInteger(r.contextWindow) ||
|
|
216
|
+
!Number.isSafeInteger(r.maxTokens) ||
|
|
217
|
+
typeof r.reasoning !== "boolean" ||
|
|
218
|
+
!Array.isArray(r.input) ||
|
|
219
|
+
!label(r.name, 256)
|
|
220
|
+
) {
|
|
221
|
+
unavailable.push({
|
|
222
|
+
profile: p.profile_name,
|
|
223
|
+
region: p.region,
|
|
224
|
+
reason: `Pi registry has no exact metadata for ${p.foundation_model_id}; update Pi to use this profile.`,
|
|
225
|
+
});
|
|
226
|
+
continue;
|
|
227
|
+
}
|
|
228
|
+
const key = slug(registryName(r.name));
|
|
229
|
+
grouped.set(key, [...(grouped.get(key) || []), p]);
|
|
230
|
+
}
|
|
231
|
+
for (const [key, rows] of grouped)
|
|
232
|
+
for (const p of rows) {
|
|
233
|
+
const r = registryEntry(p.foundation_model_id);
|
|
234
|
+
const id = rows.length === 1 ? key : `${key}--${p.region}--${p.arn.slice(p.arn.lastIndexOf("/") + 1)}`;
|
|
235
|
+
const pricing = rates.get(key);
|
|
236
|
+
const fallback =
|
|
237
|
+
r.cost &&
|
|
238
|
+
["input", "output", "cacheRead", "cacheWrite"].every(
|
|
239
|
+
(k) => Number.isFinite(r.cost[k]) && r.cost[k] >= 0,
|
|
240
|
+
)
|
|
241
|
+
? r.cost
|
|
242
|
+
: undefined;
|
|
243
|
+
if (!pricing && !fallback) {
|
|
244
|
+
unavailable.push({
|
|
245
|
+
profile: p.profile_name,
|
|
246
|
+
region: p.region,
|
|
247
|
+
reason: `No valid pricing for ${p.foundation_model_id}; update Pi or usage program models.json.`,
|
|
248
|
+
});
|
|
249
|
+
continue;
|
|
250
|
+
}
|
|
251
|
+
const promptCache = promptCacheFor(p.foundation_model_id);
|
|
252
|
+
// Only Claude profiles go on the Anthropic route; gpt, nova and the rest keep their own route.
|
|
253
|
+
const messages = api === "messages" && isAnthropicModel(p.foundation_model_id);
|
|
254
|
+
available.push({
|
|
255
|
+
id,
|
|
256
|
+
arn: p.arn,
|
|
257
|
+
foundationModelId: p.foundation_model_id,
|
|
258
|
+
aliases: [],
|
|
259
|
+
name: registryName(r.name) + (rows.length === 1 ? "" : ` (${p.region}, ${p.profile_name})`),
|
|
260
|
+
api: messages ? "anthropic-messages" : "bedrock-converse-stream",
|
|
261
|
+
baseUrl: messages
|
|
262
|
+
? `https://bedrock-runtime.${p.region}.amazonaws.com/anthropic`
|
|
263
|
+
: `https://bedrock-runtime.${p.region}.amazonaws.com`,
|
|
264
|
+
...(messages
|
|
265
|
+
? {
|
|
266
|
+
compat: {
|
|
267
|
+
...MESSAGES_COMPAT,
|
|
268
|
+
// Claude 5.5 refuses budget thinking; older Claude models still need it.
|
|
269
|
+
...(usesAdaptiveThinking(p.foundation_model_id) ? { forceAdaptiveThinking: true } : {}),
|
|
270
|
+
},
|
|
271
|
+
}
|
|
272
|
+
: {}),
|
|
273
|
+
reasoning: r.reasoning,
|
|
274
|
+
input: [...r.input],
|
|
275
|
+
contextWindow: r.contextWindow,
|
|
276
|
+
maxTokens: r.maxTokens,
|
|
277
|
+
cost: pricing || { ...fallback },
|
|
278
|
+
rateSource: pricing ? "usage program" : "Pi registry",
|
|
279
|
+
...(promptCache ? { promptCache } : {}),
|
|
280
|
+
...(thinkingLevelMap(r.thinkingLevelMap, p.foundation_model_id) ?? {}),
|
|
281
|
+
});
|
|
282
|
+
}
|
|
283
|
+
return { available, unavailable };
|
|
284
|
+
}
|
|
285
|
+
export function updateModels(
|
|
286
|
+
current: TeamBedrockModel[],
|
|
287
|
+
updated: TeamBedrockModel[],
|
|
288
|
+
): { changed: boolean; added: number; removed: number } {
|
|
289
|
+
const before = new Set(current.map((m) => `${m.id}|${m.arn}`));
|
|
290
|
+
const after = new Set(updated.map((m) => `${m.id}|${m.arn}`));
|
|
291
|
+
const changed = JSON.stringify(current) !== JSON.stringify(updated);
|
|
292
|
+
if (changed) current.splice(0, current.length, ...updated);
|
|
293
|
+
return {
|
|
294
|
+
changed,
|
|
295
|
+
added: [...after].filter((x) => !before.has(x)).length,
|
|
296
|
+
removed: [...before].filter((x) => !after.has(x)).length,
|
|
297
|
+
};
|
|
298
|
+
}
|
|
299
|
+
export function resolveSelection(models: readonly TeamBedrockModel[], selection: string): TeamBedrockModel | undefined {
|
|
300
|
+
const id = selection.replace(/^(?:amazon-bedrock|jig-bedrock)\//, "");
|
|
301
|
+
// A session saved under Pi's built-in provider names the team profile by its ARN.
|
|
302
|
+
const matches = models.filter((m) => m.id === id || m.arn === id || m.aliases.includes(id));
|
|
303
|
+
return matches.length === 1 ? matches[0] : undefined;
|
|
304
|
+
}
|
|
305
|
+
type Stream = NonNullable<ProviderConfig["streamSimple"]>;
|
|
306
|
+
export interface BedrockStreams {
|
|
307
|
+
/** Pi's Bedrock Converse stream, from its own `amazon-bedrock` provider. */
|
|
308
|
+
converse: Stream;
|
|
309
|
+
/** Pi's Anthropic Messages stream, for Bedrock's `/anthropic` route. */
|
|
310
|
+
messages?: Stream;
|
|
311
|
+
}
|
|
312
|
+
export function bedrockStream(
|
|
313
|
+
models: readonly TeamBedrockModel[],
|
|
314
|
+
streams: BedrockStreams,
|
|
315
|
+
load: () => Catalog | undefined,
|
|
316
|
+
profile?: string,
|
|
317
|
+
signedFetch?: typeof globalThis.fetch,
|
|
318
|
+
): Stream {
|
|
319
|
+
const approved = [...models];
|
|
320
|
+
return (model, context, options) => {
|
|
321
|
+
const physical = approved.find((m) => m.id === model.id);
|
|
322
|
+
if (
|
|
323
|
+
model.provider !== JIG_BEDROCK ||
|
|
324
|
+
!physical ||
|
|
325
|
+
!load()?.profiles.some(
|
|
326
|
+
(p) => p.arn === physical.arn && p.foundation_model_id === physical.foundationModelId,
|
|
327
|
+
)
|
|
328
|
+
)
|
|
329
|
+
throw new Error("Unapproved Bedrock profile; choose an available model with /model.");
|
|
330
|
+
if (
|
|
331
|
+
profile &&
|
|
332
|
+
(options?.apiKey ||
|
|
333
|
+
(options as Json)?.bearerToken ||
|
|
334
|
+
((options as Json)?.profile && (options as Json).profile !== profile) ||
|
|
335
|
+
(options?.env?.AWS_PROFILE && options.env.AWS_PROFILE !== profile))
|
|
336
|
+
)
|
|
337
|
+
throw new Error("Unapproved Bedrock authentication scope; choose the configured profile with /model.");
|
|
338
|
+
const scoped = profile
|
|
339
|
+
? {
|
|
340
|
+
...options,
|
|
341
|
+
profile,
|
|
342
|
+
env: {
|
|
343
|
+
...options?.env,
|
|
344
|
+
AWS_PROFILE: profile,
|
|
345
|
+
AWS_BEARER_TOKEN_BEDROCK: "",
|
|
346
|
+
AWS_BEDROCK_SKIP_AUTH: "0",
|
|
347
|
+
},
|
|
348
|
+
}
|
|
349
|
+
: options;
|
|
350
|
+
const wire = { ...model, ...physical, id: physical.arn, provider: "amazon-bedrock" as const };
|
|
351
|
+
if (physical.api === "anthropic-messages") {
|
|
352
|
+
const messages = streams.messages;
|
|
353
|
+
if (!messages || !signedFetch)
|
|
354
|
+
throw new Error("The Bedrock Messages route is not available; reload Pi and check the Jig install.");
|
|
355
|
+
const previous = (options as Json)?.onPayload as
|
|
356
|
+
| ((payload: unknown, selected: unknown) => unknown)
|
|
357
|
+
| undefined;
|
|
358
|
+
return messages({ ...wire, api: "anthropic-messages" }, context, {
|
|
359
|
+
...scoped,
|
|
360
|
+
// Pi's Anthropic client wants some key before it builds a request. It never reaches Bedrock: the signed
|
|
361
|
+
// fetch removes both this and any Authorization header before signing.
|
|
362
|
+
apiKey: "bedrock-sigv4",
|
|
363
|
+
fetch: signedFetch,
|
|
364
|
+
// Sonnet 5.5 refuses Pi's thinking-off shape; every other payload is left exactly as it came, and Pi's own
|
|
365
|
+
// payload hook still runs first.
|
|
366
|
+
onPayload: async (payload: unknown, selected: unknown) => {
|
|
367
|
+
const next = previous === undefined ? payload : await previous(payload, selected);
|
|
368
|
+
const value = (next ?? payload) as Record<string, unknown>;
|
|
369
|
+
return rewriteThinkingOff(value, physical.foundationModelId);
|
|
370
|
+
},
|
|
371
|
+
} as typeof options);
|
|
372
|
+
}
|
|
373
|
+
return streams.converse({ ...wire, api: "bedrock-converse-stream" }, context, scoped);
|
|
374
|
+
};
|
|
375
|
+
}
|
|
376
|
+
export async function registerTeamBedrock(
|
|
377
|
+
pi: Pick<ExtensionAPI, "registerProvider">,
|
|
378
|
+
_agentDir: string,
|
|
379
|
+
native?: import("@earendil-works/pi-ai").Provider,
|
|
380
|
+
models: TeamBedrockModel[] = [],
|
|
381
|
+
load: () => Catalog | undefined = () => undefined,
|
|
382
|
+
profile?: string,
|
|
383
|
+
messages?: { stream: Stream; streamSimple: Stream },
|
|
384
|
+
signedFetch?: typeof globalThis.fetch,
|
|
385
|
+
): Promise<TeamBedrockModel[]> {
|
|
386
|
+
if (!native) {
|
|
387
|
+
const { builtinProviders } = await import("@earendil-works/pi-ai/providers/all");
|
|
388
|
+
native = builtinProviders().find((provider) => provider.id === "amazon-bedrock");
|
|
389
|
+
if (!native) throw new Error("Pi Bedrock provider unavailable; update Pi and reload.");
|
|
390
|
+
}
|
|
391
|
+
const picker = models.map(
|
|
392
|
+
({ arn: _arn, aliases: _aliases, foundationModelId: _foundation, rateSource: _rate, ...m }) => ({
|
|
393
|
+
...m,
|
|
394
|
+
type: "chat" as const,
|
|
395
|
+
provider: JIG_BEDROCK,
|
|
396
|
+
}),
|
|
397
|
+
);
|
|
398
|
+
pi.registerProvider({
|
|
399
|
+
id: JIG_BEDROCK,
|
|
400
|
+
name: "Bedrock (Jig)",
|
|
401
|
+
auth: profile
|
|
402
|
+
? {
|
|
403
|
+
...native.auth,
|
|
404
|
+
apiKey: {
|
|
405
|
+
...native.auth?.apiKey,
|
|
406
|
+
name: native.auth?.apiKey?.name ?? "AWS profile",
|
|
407
|
+
resolve: async () => ({
|
|
408
|
+
auth: {},
|
|
409
|
+
env: { AWS_PROFILE: profile },
|
|
410
|
+
source: "configured AWS profile",
|
|
411
|
+
}),
|
|
412
|
+
},
|
|
413
|
+
}
|
|
414
|
+
: native.auth,
|
|
415
|
+
getModels: () => picker,
|
|
416
|
+
getAllModels: () => picker,
|
|
417
|
+
stream: bedrockStream(
|
|
418
|
+
models,
|
|
419
|
+
{ converse: native.stream as Stream, ...(messages ? { messages: messages.stream } : {}) },
|
|
420
|
+
load,
|
|
421
|
+
profile,
|
|
422
|
+
signedFetch,
|
|
423
|
+
) as typeof native.stream,
|
|
424
|
+
streamSimple: bedrockStream(
|
|
425
|
+
models,
|
|
426
|
+
{
|
|
427
|
+
converse: native.streamSimple,
|
|
428
|
+
...(messages ? { messages: messages.streamSimple } : {}),
|
|
429
|
+
},
|
|
430
|
+
load,
|
|
431
|
+
profile,
|
|
432
|
+
signedFetch,
|
|
433
|
+
),
|
|
434
|
+
});
|
|
435
|
+
pi.registerProvider("amazon-bedrock", { models: [] });
|
|
436
|
+
return models;
|
|
437
|
+
}
|
|
438
|
+
export function registerBedrockSelectionGuard(
|
|
439
|
+
pi: Pick<ExtensionAPI, "on" | "setModel">,
|
|
440
|
+
models: TeamBedrockModel[],
|
|
441
|
+
options: { agentDir?: string; explicitSelection?: () => boolean } = {},
|
|
442
|
+
): { blockUntilChoice(): void } {
|
|
443
|
+
let unresolvedSavedSelection = false;
|
|
444
|
+
const warning =
|
|
445
|
+
"Saved Bedrock model is unavailable or ambiguous for this team; choose a model with /model before sending a request.";
|
|
446
|
+
const unsupported = (ctx: ExtensionContext) =>
|
|
447
|
+
unresolvedSavedSelection ||
|
|
448
|
+
(isJigBedrock(ctx.model) &&
|
|
449
|
+
(ctx.model?.provider !== JIG_BEDROCK || !models.some((m) => m.id === ctx.model?.id)));
|
|
450
|
+
const restore = async (event: { type?: string; reason?: string }, ctx: ExtensionContext) => {
|
|
451
|
+
// Pi 1.2 reports a session switch as session_start with reason new, resume or fork.
|
|
452
|
+
const switched =
|
|
453
|
+
event.type === "session_start" &&
|
|
454
|
+
(event.reason === "new" || event.reason === "resume" || event.reason === "fork");
|
|
455
|
+
let warned = false;
|
|
456
|
+
const selected = [...ctx.sessionManager.getBranch()]
|
|
457
|
+
.reverse()
|
|
458
|
+
.find(
|
|
459
|
+
(e: Json) => e.type === "model_change" || (e.type === "message" && e.message?.role === "assistant"),
|
|
460
|
+
) as Json;
|
|
461
|
+
let provider = selected?.type === "model_change" ? selected.provider : selected?.message?.provider;
|
|
462
|
+
let id = selected?.type === "model_change" ? selected.modelId : selected?.message?.model;
|
|
463
|
+
if (!selected && !options.explicitSelection?.()) {
|
|
464
|
+
try {
|
|
465
|
+
const settings = JSON.parse(
|
|
466
|
+
readFileSync(
|
|
467
|
+
join(
|
|
468
|
+
options.agentDir || process.env.PI_CODING_AGENT_DIR || join(homedir(), ".pi/agent"),
|
|
469
|
+
"settings.json",
|
|
470
|
+
),
|
|
471
|
+
"utf8",
|
|
472
|
+
),
|
|
473
|
+
);
|
|
474
|
+
provider = settings.defaultProvider;
|
|
475
|
+
id = settings.defaultModel;
|
|
476
|
+
} catch {
|
|
477
|
+
/* No saved default. */
|
|
478
|
+
}
|
|
479
|
+
}
|
|
480
|
+
if ((!options.explicitSelection?.() || switched) && isJigBedrock({ provider })) {
|
|
481
|
+
const match = resolveSelection(models, id || ""),
|
|
482
|
+
target = match && ctx.modelRegistry.find(JIG_BEDROCK, match.id);
|
|
483
|
+
if (!target) {
|
|
484
|
+
unresolvedSavedSelection = true;
|
|
485
|
+
ctx.ui.notify(warning, "warning");
|
|
486
|
+
warned = true;
|
|
487
|
+
} else if (ctx.model?.provider !== target.provider || ctx.model.id !== target.id) {
|
|
488
|
+
if (!(await pi.setModel(target))) {
|
|
489
|
+
unresolvedSavedSelection = true;
|
|
490
|
+
ctx.ui.notify(warning, "warning");
|
|
491
|
+
warned = true;
|
|
492
|
+
} else unresolvedSavedSelection = false;
|
|
493
|
+
}
|
|
494
|
+
}
|
|
495
|
+
if (!warned && unsupported(ctx)) ctx.ui.notify(warning, "warning");
|
|
496
|
+
};
|
|
497
|
+
pi.on("session_start", restore);
|
|
498
|
+
pi.on("model_select", (event, ctx) => {
|
|
499
|
+
if (event.source === "set" || event.source === "cycle") unresolvedSavedSelection = false;
|
|
500
|
+
if (unsupported(ctx)) ctx.ui.notify(warning, "warning");
|
|
501
|
+
});
|
|
502
|
+
pi.on("input", (event, ctx) => {
|
|
503
|
+
if (!unsupported(ctx)) return { action: "continue" };
|
|
504
|
+
if (ctx.mode === "tui" && event.source === "interactive" && !event.streamingBehavior)
|
|
505
|
+
ctx.ui.setEditorText(event.text);
|
|
506
|
+
ctx.ui.notify(warning, "warning");
|
|
507
|
+
return { action: "handled" };
|
|
508
|
+
});
|
|
509
|
+
return {
|
|
510
|
+
blockUntilChoice: () => {
|
|
511
|
+
unresolvedSavedSelection = true;
|
|
512
|
+
},
|
|
513
|
+
};
|
|
514
|
+
}
|
package/src/check.ts
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
// The two small checks the `bedrock` section shares. Errors name the key and never repeat a value.
|
|
2
|
+
|
|
3
|
+
export type Check = (raw: unknown, errors: string[]) => void;
|
|
4
|
+
|
|
5
|
+
export function isObject(value: unknown): value is Record<string, unknown> {
|
|
6
|
+
return typeof value === "object" && value !== null && !Array.isArray(value);
|
|
7
|
+
}
|
|
8
|
+
|
|
9
|
+
/** Reports every key of `raw` that is not in `known`. */
|
|
10
|
+
export function knownKeys(raw: Record<string, unknown>, known: readonly string[], at: string, errors: string[]): void {
|
|
11
|
+
for (const key of Object.keys(raw)) {
|
|
12
|
+
if (!known.includes(key)) errors.push(`${at}.${key.slice(0, 40)} is not a known field`);
|
|
13
|
+
}
|
|
14
|
+
}
|
package/src/child.ts
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* A subagent child is a short-lived Pi the parent started: the subagents extension sets PI_SUBAGENT_ID in it. PI_SUBAGENT_CHILD=1
|
|
3
|
+
* is the older name, still accepted. Children skip display work and background AWS and usage program calls: nobody reads a footer there
|
|
4
|
+
* and the parent already keeps the catalog and the usage figures fresh.
|
|
5
|
+
*/
|
|
6
|
+
export const isSubagentChild = (env: NodeJS.ProcessEnv = process.env): boolean =>
|
|
7
|
+
!!env.PI_SUBAGENT_ID?.trim() || env.PI_SUBAGENT_CHILD === "1";
|