Repository navigation
fix(bedrock): non-empty Converse streams in the release build and cached tokens billed once #1425
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
anandgupta42
wants to merge
8
commits into
main
Choose a base branch
from
fix/bedrock-stream-and-cost
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
8 commits
Select commit
Hold shift + click to select a range
a23e2bd
fix(bedrock): keep Converse streams non-empty in the release build
05bced1
fix(session): bill cached input tokens once for Anthropic and Bedrock
5b02f21
fix(learn): pass the SDK's uncached input count to session cost accou…
02d507b
fix(session): count cache writes in the long-context price tier and r…
b23e393
fix(session): keep detail-only cache writes and reads consistent acro…
aaaf42f
fix(session): trust detail cache counts only alongside `noCacheTokens`
5249f8a
fix(session): keep detail cache fallbacks for inclusive providers and…
94df9bc
fix(session): omit `tokens_cache_read` when accounting rejected the d…
anandgupta42 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| // Bun.build plugin shared by script/build.ts and its regression test. | ||
| // | ||
| // The release build bundles with `conditions: ["browser"]` (needed by the TUI | ||
| // dependencies). Under that condition `@smithy/core/serde` resolves to | ||
| // `index.browser.js`, where the Node-only exports `fromString` and | ||
| // `fromArrayBuffer` are placeholder symbols (`Symbol.for("node-only")`). | ||
| // `@smithy/util-buffer-from` re-exports exactly those two names and | ||
| // `@smithy/util-utf8` calls them while `@ai-sdk/amazon-bedrock` decodes the | ||
| // `application/vnd.amazon.eventstream` body. The call throws a TypeError, which | ||
| // the SDK's decoder swallows, so every Bedrock Converse stream ended with zero | ||
| // events: finish "other", no text, no usage, no error. | ||
| // | ||
| // This plugin redirects only that one import (`@smithy/core/serde` requested | ||
| // from inside `@smithy/util-buffer-from`) to the package's Node build, which | ||
| // is the file a source run already uses. Every other importer of | ||
| // `@smithy/core/serde` and every other package keeps its resolution. | ||
| import path from "node:path" | ||
| import type { BunPlugin } from "bun" | ||
|
|
||
| const IMPORTER = /[\\/]@smithy[\\/]util-buffer-from[\\/]/ | ||
|
|
||
| export function smithyNodeSerdePlugin(): BunPlugin { | ||
| return { | ||
| name: "smithy-util-buffer-from-node-serde", | ||
| setup(build) { | ||
| build.onResolve({ filter: /^@smithy\/core\/serde$/ }, (args) => { | ||
| if (!IMPORTER.test(args.importer)) return undefined | ||
| // Resolve from the importer's directory with Bun's default (non-browser) conditions. | ||
| try { | ||
| return { path: Bun.resolveSync("@smithy/core/serde", path.dirname(args.importer)) } | ||
| } catch (cause) { | ||
| throw new Error( | ||
| `smithy-node-serde plugin: cannot resolve @smithy/core/serde from ${args.importer}; Bedrock streaming would be empty in the built binary`, | ||
| { cause }, | ||
| ) | ||
| } | ||
| }) | ||
| }, | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
113 changes: 113 additions & 0 deletions
113
packages/opencode/test/script/bedrock-eventstream-build.test.ts
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,113 @@ | ||
| import { describe, test, expect } from "bun:test" | ||
| import fs from "node:fs" | ||
| import os from "node:os" | ||
| import path from "node:path" | ||
| import { smithyNodeSerdePlugin } from "../../script/smithy-node-serde-plugin" | ||
|
|
||
| /** | ||
| * Regression: Bedrock Converse streams came back empty (finish "other", no text, no usage, no error) | ||
| * in the built binary. script/build.ts bundles with `conditions: ["browser"]` for a Bun target; under | ||
| * that combination `@smithy/core/serde` exports placeholder `fromString`/`fromArrayBuffer`, the event | ||
| * stream decoder throws on its first frame and the SDK swallows the error. | ||
| * | ||
| * These tests bundle a probe that decodes a synthetic `application/vnd.amazon.eventstream` body through | ||
| * the real @ai-sdk/amazon-bedrock provider, using the same `conditions` and a Bun target as the release | ||
| * build, then run the bundle. No network or credentials are involved. | ||
| */ | ||
|
|
||
| const bundler = path.join(import.meta.dir, "bedrock-eventstream-bundle.ts") | ||
|
|
||
| async function bundleAndRun(mode: "plugin" | "none") { | ||
| const outdir = fs.mkdtempSync(path.join(os.tmpdir(), "bedrock-eventstream-")) | ||
| try { | ||
| const build = Bun.spawn([process.execPath, bundler, outdir, mode], { stdout: "pipe", stderr: "pipe" }) | ||
| const buildErr = await new Response(build.stderr).text() | ||
| expect(await build.exited, buildErr).toBe(0) | ||
| const run = Bun.spawn([process.execPath, path.join(outdir, "bedrock-eventstream-probe.js")], { | ||
| stdout: "pipe", | ||
| stderr: "pipe", | ||
| }) | ||
| const [out, err] = await Promise.all([new Response(run.stdout).text(), new Response(run.stderr).text()]) | ||
| expect(await run.exited, err).toBe(0) | ||
| return JSON.parse(out.trim().split("\n").pop()!) | ||
| } finally { | ||
| fs.rmSync(outdir, { recursive: true, force: true }) | ||
| } | ||
| } | ||
|
|
||
| describe("Bedrock event-stream decoding in the release build's resolution conditions", () => { | ||
| test("with the build plugin the stream yields text, finish reason and token counts", async () => { | ||
| const result = await bundleAndRun("plugin") | ||
| expect(result).toEqual({ text: "hello", finish: { unified: "stop", raw: "end_turn" }, input: 4, output: 2 }) | ||
| }, 60_000) | ||
|
|
||
| // Control: proves the probe detects the defect. If an @smithy/core or @smithy/util-buffer-from upgrade | ||
| // makes this fail (the stream is no longer empty), the plugin is no longer needed and can be removed | ||
| // together with this test. | ||
| test("without the plugin the same bundle is empty (control for the probe)", async () => { | ||
| const result = await bundleAndRun("none") | ||
| expect(result.text).toBe("") | ||
| expect(result.finish.unified).toBe("other") | ||
| }, 60_000) | ||
| }) | ||
|
|
||
| describe("smithyNodeSerdePlugin scope", () => { | ||
| function resolver() { | ||
| let cb!: (args: { importer: string }) => unknown | ||
| smithyNodeSerdePlugin().setup({ onResolve: (_o: unknown, f: unknown) => (cb = f as typeof cb) } as any) | ||
| return cb | ||
| } | ||
|
|
||
| test("redirects @smithy/core/serde only for @smithy/util-buffer-from importers", () => { | ||
| const resolve = resolver() | ||
| const bufferFrom = require.resolve("@smithy/util-buffer-from/package.json", { | ||
| paths: [ | ||
| path.dirname( | ||
| require.resolve("@smithy/util-utf8/package.json", { | ||
| paths: [path.dirname(require.resolve("@ai-sdk/amazon-bedrock/package.json"))], | ||
| }), | ||
| ), | ||
| ], | ||
| }) | ||
| const hit = resolve({ importer: path.join(path.dirname(bufferFrom), "dist-es", "index.js") }) as { path: string } | ||
| expect(hit.path).toContain("serde") | ||
| expect(hit.path).not.toContain("browser") | ||
| }) | ||
|
|
||
| test("fails with an actionable error (and keeps the cause) when the serde module cannot be resolved", () => { | ||
| const resolve = resolver() | ||
| const original = Bun.resolveSync | ||
| const boom = new Error("not found") | ||
| ;(Bun as any).resolveSync = () => { | ||
| throw boom | ||
| } | ||
| let caught: any | ||
| try { | ||
| resolve({ importer: "/x/node_modules/@smithy/util-buffer-from/dist-es/index.js" }) | ||
| } catch (err) { | ||
| caught = err | ||
| } finally { | ||
| ;(Bun as any).resolveSync = original | ||
| } | ||
| expect(caught).toBeDefined() | ||
| expect(caught.message).toContain("cannot resolve @smithy/core/serde") | ||
| expect(caught.message).toContain("/x/node_modules/@smithy/util-buffer-from/dist-es/index.js") | ||
| expect(caught.cause).toBe(boom) | ||
| }) | ||
|
|
||
| test("leaves every other importer untouched", () => { | ||
| const resolve = resolver() | ||
| expect(resolve({ importer: "/x/node_modules/@smithy/util-utf8/dist-es/fromUtf8.js" })).toBeUndefined() | ||
| expect(resolve({ importer: "/x/node_modules/@smithy/eventstream-codec/dist-es/index.js" })).toBeUndefined() | ||
| expect(resolve({ importer: "/x/src/index.ts" })).toBeUndefined() | ||
| }) | ||
| }) | ||
|
|
||
| describe("script/build.ts wiring", () => { | ||
| test("the release Bun.build call registers smithyNodeSerdePlugin alongside the browser condition", () => { | ||
| const src = fs.readFileSync(path.join(import.meta.dir, "../../script/build.ts"), "utf8") | ||
| const call = src.slice(src.indexOf("await Bun.build({")) | ||
| expect(call).toContain('conditions: ["browser"]') | ||
| expect(call.slice(0, call.indexOf("entrypoints:"))).toContain("smithyNodeSerdePlugin()") | ||
| }) | ||
| }) |
20 changes: 20 additions & 0 deletions
20
packages/opencode/test/script/bedrock-eventstream-bundle.ts
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| // Usage: bun bedrock-eventstream-bundle.ts <outdir> <plugin|none> | ||
| // Bundles the probe with the release build's resolution conditions. Runs in its own process because | ||
| // Bun.build does not tolerate a second differently-configured build over the same graph in-process. | ||
| import path from "node:path" | ||
| import { smithyNodeSerdePlugin } from "../../script/smithy-node-serde-plugin" | ||
|
|
||
| const [outdir, mode] = process.argv.slice(2) | ||
| const root = path.resolve(import.meta.dir, "../..") | ||
| const built = await Bun.build({ | ||
| entrypoints: [path.join(import.meta.dir, "bedrock-eventstream-probe.ts")], | ||
| outdir, | ||
| conditions: ["browser"], | ||
| target: "bun", | ||
| tsconfig: path.join(root, "tsconfig.json"), | ||
| plugins: mode === "plugin" ? [smithyNodeSerdePlugin()] : [], | ||
| }) | ||
| if (!built.success) { | ||
| console.error(built.logs.join("\n")) | ||
| process.exit(1) | ||
| } |
38 changes: 38 additions & 0 deletions
38
packages/opencode/test/script/bedrock-eventstream-frames.ts
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,38 @@ | ||
| const enc = new TextEncoder() | ||
|
|
||
| function header(name: string, value: string): Uint8Array { | ||
| const n = enc.encode(name) | ||
| const v = enc.encode(value) | ||
| const out = new Uint8Array(1 + n.length + 1 + 2 + v.length) | ||
| out[0] = n.length | ||
| out.set(n, 1) | ||
| out[1 + n.length] = 7 // header value type: string | ||
| new DataView(out.buffer).setUint16(2 + n.length, v.length, false) | ||
| out.set(v, 4 + n.length) | ||
| return out | ||
| } | ||
|
|
||
| /** Encode one `application/vnd.amazon.eventstream` message frame carrying a JSON payload. */ | ||
| export function eventStreamFrame(eventType: string, payload: unknown): Uint8Array { | ||
| const headers = [ | ||
| header(":event-type", eventType), | ||
| header(":content-type", "application/json"), | ||
| header(":message-type", "event"), | ||
| ] | ||
| const headersLen = headers.reduce((a, h) => a + h.length, 0) | ||
| const body = enc.encode(JSON.stringify(payload)) | ||
| const total = 12 + headersLen + body.length + 4 | ||
| const out = new Uint8Array(total) | ||
| const dv = new DataView(out.buffer) | ||
| dv.setUint32(0, total, false) | ||
| dv.setUint32(4, headersLen, false) | ||
| dv.setUint32(8, Bun.hash.crc32(out.subarray(0, 8)) as number, false) | ||
| let off = 12 | ||
| for (const h of headers) { | ||
| out.set(h, off) | ||
| off += h.length | ||
| } | ||
| out.set(body, off) | ||
| dv.setUint32(total - 4, Bun.hash.crc32(out.subarray(0, total - 4)) as number, false) | ||
| return out | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.