Skip to content

Commit 41bd1a1

Browse files
nborges-awsnborges-aws
committed
fix: remove first-response behavior, sanitize URLs, and verify draft state when writing
1 parent b6c4cb8 commit 41bd1a1

35 files changed

Lines changed: 474 additions & 87 deletions

‎src/core/datasetUpdate.test.ts‎

Lines changed: 66 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -256,11 +256,7 @@ describe("EvalClient.updateDatasetExamples", () => {
256256

257257
expect(addCommands).toHaveLength(2);
258258
for (const command of addCommands) {
259-
const requestBody = {
260-
source: command.input.source,
261-
clientToken: command.input.clientToken,
262-
};
263-
expect(Buffer.byteLength(JSON.stringify(requestBody), "utf8")).toBeLessThanOrEqual(
259+
expect(Buffer.byteLength(JSON.stringify(command.input), "utf8")).toBeLessThanOrEqual(
264260
PAYLOAD_LIMIT_BYTES,
265261
);
266262
}
@@ -292,12 +288,27 @@ describe("EvalClient.updateDatasetExamples", () => {
292288

293289
expect(updateCommands).toHaveLength(2);
294290
expect(updateCommands.every((command) => command.input.examples?.length === 1)).toBe(true);
291+
expect(
292+
updateCommands.every(
293+
(command) =>
294+
Buffer.byteLength(JSON.stringify(command.input), "utf8") <= PAYLOAD_LIMIT_BYTES,
295+
),
296+
).toBe(true);
295297
});
296298

297-
test("rejects an individually oversized example before mutating the dataset", async () => {
298-
const localPath = tempFile(
299-
jsonl({ scenario_id: "too-large", value: "x".repeat(PAYLOAD_LIMIT_BYTES) }),
300-
);
299+
test("includes datasetId when rejecting an oversized mutation before changing the dataset", async () => {
300+
const clientToken = "0".repeat(36);
301+
const example = { scenario_id: "too-large", value: "" };
302+
const bodyWithoutDatasetId = {
303+
source: { inlineExamples: { examples: [example] } },
304+
clientToken,
305+
};
306+
const valueBytes =
307+
PAYLOAD_LIMIT_BYTES - Buffer.byteLength(JSON.stringify(bodyWithoutDatasetId));
308+
example.value = "x".repeat(valueBytes);
309+
expect(Buffer.byteLength(JSON.stringify(bodyWithoutDatasetId))).toBe(PAYLOAD_LIMIT_BYTES);
310+
311+
const localPath = tempFile(jsonl(example));
301312
const commands: unknown[] = [];
302313
const fetch = (() => {
303314
throw new Error("fetch should not be called");
@@ -497,6 +508,52 @@ describe("EvalClient.updateDatasetExamples", () => {
497508
expect(checkpoint.filter((row) => row.exampleId !== undefined)).toHaveLength(1000);
498509
});
499510

511+
test("preserves concurrent edits and stops after writing assigned IDs to a recovery file", async () => {
512+
const localRows = Array.from({ length: 1001 }, (_, i) => ({ scenario_id: `new-${i}` }));
513+
const localPath = tempFile(jsonl(...localRows));
514+
const editedContents = jsonl({ scenario_id: "edited-while-request-was-running" });
515+
const commands: unknown[] = [];
516+
let addCalls = 0;
517+
const client = new EvalClient(
518+
stubClients({
519+
commands,
520+
dataset: { datasetId: "d-1", datasetVersion: "DRAFT", status: "ACTIVE", exampleCount: 0 },
521+
addIds: Array.from({ length: 1000 }, (_, i) => `fresh-${i}`),
522+
beforeSend: (command) => {
523+
if (!(command instanceof AddDatasetExamplesCommand) || ++addCalls !== 1) return;
524+
writeFileSync(localPath, editedContents);
525+
},
526+
}),
527+
(() => {
528+
throw new Error("fetch should not be called");
529+
}) as unknown as CoreFetch,
530+
);
531+
532+
let conflict: unknown;
533+
try {
534+
await client.updateDatasetExamples("d-1", localPath, OPTIONS);
535+
} catch (error) {
536+
conflict = error;
537+
}
538+
539+
expect(conflict).toBeInstanceOf(InputValidationError);
540+
const recoveryFilePath = (conflict as InputValidationError).meta.recoveryFilePath;
541+
expect(recoveryFilePath).toBeString();
542+
expect(readFileSync(localPath, "utf8")).toBe(editedContents);
543+
544+
const recovered = readFileSync(recoveryFilePath as string, "utf8")
545+
.trim()
546+
.split("\n")
547+
.map((line) => JSON.parse(line));
548+
expect(recovered[0]?.exampleId).toBe("fresh-0");
549+
expect(recovered[999]?.exampleId).toBe("fresh-999");
550+
expect(recovered[1000]).not.toHaveProperty("exampleId");
551+
expect(commands.filter((command) => command instanceof AddDatasetExamplesCommand)).toHaveLength(
552+
1,
553+
);
554+
expect(commands.at(-1)).toBeInstanceOf(GetDatasetCommand);
555+
});
556+
500557
test("aborts a pending mutation through the SDK request options", async () => {
501558
const localPath = tempFile(jsonl({ scenario_id: "new" }));
502559
const commands: unknown[] = [];

‎src/core/eval.tsx‎

Lines changed: 56 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ import {
5252
type CloudWatchFilterConfig,
5353
} from "@aws-sdk/client-bedrock-agentcore";
5454
import { randomUUID } from "node:crypto";
55+
import { basename, dirname, extname, join } from "node:path";
5556
import { Transform } from "node:stream";
5657
import { setTimeout as sleep } from "node:timers/promises";
5758
import {
@@ -739,17 +740,18 @@ export class EvalClient implements CoreEvalClient {
739740
const deleteBatches = buildDatasetExampleBatches({
740741
items: diff.deleteIds,
741742
payloadItem: (exampleId) => exampleId,
742-
requestBody: (exampleIds, clientToken) => ({ exampleIds, clientToken }),
743+
requestBody: (exampleIds, clientToken) => ({ datasetId: id, exampleIds, clientToken }),
743744
});
744745
const updateBatches = buildDatasetExampleBatches({
745746
items: diff.updates,
746747
payloadItem: (example) => example,
747-
requestBody: (examples, clientToken) => ({ examples, clientToken }),
748+
requestBody: (examples, clientToken) => ({ datasetId: id, examples, clientToken }),
748749
});
749750
const additionBatches = buildDatasetExampleBatches({
750751
items: diff.additions,
751752
payloadItem: (addition) => addition.content,
752753
requestBody: (examples, clientToken) => ({
754+
datasetId: id,
753755
source: { inlineExamples: { examples } },
754756
clientToken,
755757
}),
@@ -791,6 +793,8 @@ export class EvalClient implements CoreEvalClient {
791793

792794
const completedAdditions: Addition[] = [];
793795
const assignedIds: string[] = [];
796+
let expectedLocalText = localText;
797+
let checkpointConflict: InputValidationError | undefined;
794798
await runDatasetExampleBatches({
795799
batches: additionBatches,
796800
datasetId: id,
@@ -812,17 +816,46 @@ export class EvalClient implements CoreEvalClient {
812816
completedAdditions.push(...additions);
813817
assignedIds.push(...(response.exampleIds ?? []));
814818
const nextLocalText = applyExampleIds(localExamples, completedAdditions, assignedIds);
819+
820+
const currentLocalText = await readTextFile(filePath);
821+
if (currentLocalText !== expectedLocalText) {
822+
const recoveryFilePath = datasetRecoveryFilePath(filePath);
823+
try {
824+
await atomicWrite(recoveryFilePath, nextLocalText);
825+
} catch (error) {
826+
throw new FileWriteError(
827+
`Dataset "${id}" changed while the update was running and its recovered IDs ` +
828+
`could not be written to ${recoveryFilePath}`,
829+
{
830+
cause: error,
831+
meta: { datasetId: id, filePath, recoveryFilePath },
832+
},
833+
);
834+
}
835+
checkpointConflict = new InputValidationError(
836+
`Dataset file "${filePath}" changed while the update was running. ` +
837+
`The file was left untouched and the reconciled content was written to ` +
838+
`"${recoveryFilePath}".`,
839+
{ meta: { datasetId: id, filePath, recoveryFilePath } },
840+
);
841+
return;
842+
}
843+
815844
try {
816845
// The remote request has already succeeded, so checkpoint its IDs even
817846
// if cancellation arrives before the next poll or batch.
818847
await atomicWrite(filePath, nextLocalText);
848+
expectedLocalText = nextLocalText;
819849
} catch (error) {
820850
throw new FileWriteError(`Could not write dataset "${id}" to ${filePath}`, {
821851
cause: error,
822852
meta: { datasetId: id, filePath },
823853
});
824854
}
825855
},
856+
afterBatchSettled: async (): Promise<void> => {
857+
if (checkpointConflict) throw checkpointConflict;
858+
},
826859
});
827860

828861
return {
@@ -897,6 +930,15 @@ async function readLocalDatasetFile(filePath: string, signal?: AbortSignal): Pro
897930
}
898931
}
899932

933+
function datasetRecoveryFilePath(filePath: string): string {
934+
const extension = extname(filePath);
935+
const stem = basename(filePath, extension);
936+
return join(
937+
dirname(filePath),
938+
`${stem}.agentcore-recovery-${randomUUID()}${extension || ".jsonl"}`,
939+
);
940+
}
941+
900942
function parseRemoteDatasetExamples(text: string): ReturnType<typeof parseJsonl> {
901943
try {
902944
return parseJsonl(text, "remote dataset");
@@ -982,14 +1024,25 @@ async function runDatasetExampleBatches<T, R>(options: {
9821024
onBatchStart?: () => void;
9831025
operation: (batch: T[], clientToken: string) => Promise<R>;
9841026
afterOperation?: (batch: T[], response: R) => Promise<void>;
1027+
afterBatchSettled?: (batch: T[], response: R) => Promise<void>;
9851028
}): Promise<void> {
986-
const { batches, datasetId, control, signal, onBatchStart, operation, afterOperation } = options;
1029+
const {
1030+
batches,
1031+
datasetId,
1032+
control,
1033+
signal,
1034+
onBatchStart,
1035+
operation,
1036+
afterOperation,
1037+
afterBatchSettled,
1038+
} = options;
9871039
for (const batch of batches) {
9881040
signal?.throwIfAborted();
9891041
onBatchStart?.();
9901042
const response = await operation(batch.items, batch.clientToken);
9911043
await afterOperation?.(batch.items, response);
9921044
await waitForDatasetActive(control, datasetId, signal);
1045+
await afterBatchSettled?.(batch.items, response);
9931046
}
9941047
}
9951048

‎src/handlers/eval/dataset/__fixtures__/CreateDatasetCommand.7dac68e9cb7b945b.json‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
{
2-
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-jzVpQaA5It",
3-
"datasetId": "agentcore_cli_dataset_fixture-jzVpQaA5It",
2+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-qhP7ipAkaf",
3+
"datasetId": "agentcore_cli_dataset_fixture-qhP7ipAkaf",
44
"status": "CREATING",
55
"createdAt": {
6-
"$date": "2026-08-07T19:02:41.330Z"
6+
"$date": "2026-08-13T14:21:59.997Z"
77
}
88
}

‎src/handlers/eval/dataset/__fixtures__/CreateDatasetCommand.8f61bc771fe0e132.json‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
{
2-
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture_second-t9vDWy80mc",
3-
"datasetId": "agentcore_cli_dataset_fixture_second-t9vDWy80mc",
2+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture_second-pXAm28GIRv",
3+
"datasetId": "agentcore_cli_dataset_fixture_second-pXAm28GIRv",
44
"status": "CREATING",
55
"createdAt": {
6-
"$date": "2026-08-07T19:03:01.523Z"
6+
"$date": "2026-08-13T14:22:19.871Z"
77
}
88
}

‎src/handlers/eval/dataset/__fixtures__/CreateDatasetVersionCommand.a0f31edf2ef45cbc.json‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
{
2+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-qhP7ipAkaf",
3+
"datasetId": "agentcore_cli_dataset_fixture-qhP7ipAkaf",
4+
"status": "UPDATING",
5+
"datasetVersion": "1",
6+
"createdAt": {
7+
"$date": "2026-08-13T14:21:59.997Z"
8+
}
9+
}

‎src/handlers/eval/dataset/__fixtures__/DeleteDatasetCommand.47a08dc8d6da4a3f.json‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
{
2+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-qhP7ipAkaf",
3+
"datasetId": "agentcore_cli_dataset_fixture-qhP7ipAkaf",
4+
"status": "UPDATING",
5+
"datasetVersion": "1",
6+
"updatedAt": {
7+
"$date": "2026-08-13T14:22:36.056Z"
8+
}
9+
}

‎src/handlers/eval/dataset/__fixtures__/DeleteDatasetCommand.a0f31edf2ef45cbc.json‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
{
2+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-qhP7ipAkaf",
3+
"datasetId": "agentcore_cli_dataset_fixture-qhP7ipAkaf",
4+
"status": "DELETING",
5+
"datasetVersion": "DRAFT",
6+
"updatedAt": {
7+
"$date": "2026-08-13T14:22:41.236Z"
8+
}
9+
}

‎src/handlers/eval/dataset/__fixtures__/GetDatasetCommand.47a08dc8d6da4a3f.json‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
{
2+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-qhP7ipAkaf",
3+
"datasetId": "agentcore_cli_dataset_fixture-qhP7ipAkaf",
4+
"datasetVersion": "1",
5+
"datasetName": "agentcore_cli_dataset_fixture",
6+
"status": "ACTIVE",
7+
"schemaType": "AGENTCORE_EVALUATION_PREDEFINED_V1",
8+
"exampleCount": 2,
9+
"createdAt": {
10+
"$date": "2026-08-13T14:21:59.997Z"
11+
},
12+
"updatedAt": {
13+
"$date": "2026-08-13T14:22:30.931Z"
14+
},
15+
"description": "Recorded fixture for the dataset commands",
16+
"downloadUrl": "https://agentcoredatasets685197708687-284077270265-us-west-2-an.s3.us-west-2.amazonaws.com/685197708687/datasets/agentcore_cli_dataset_fixture-qhP7ipAkaf/versions/0000000001/dataset.jsonl",
17+
"downloadUrlExpiresAt": {
18+
"$date": "2026-08-13T14:27:35.877Z"
19+
}
20+
}

‎src/handlers/eval/dataset/__fixtures__/GetDatasetCommand.a0f31edf2ef45cbc.json‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
{
2+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-qhP7ipAkaf",
3+
"datasetId": "agentcore_cli_dataset_fixture-qhP7ipAkaf",
4+
"datasetVersion": "DRAFT",
5+
"datasetName": "agentcore_cli_dataset_fixture",
6+
"status": "ACTIVE",
7+
"schemaType": "AGENTCORE_EVALUATION_PREDEFINED_V1",
8+
"exampleCount": 2,
9+
"createdAt": {
10+
"$date": "2026-08-13T14:21:59.997Z"
11+
},
12+
"updatedAt": {
13+
"$date": "2026-08-13T14:22:01.465Z"
14+
},
15+
"description": "Recorded fixture for the dataset commands",
16+
"draftStatus": "MODIFIED",
17+
"downloadUrl": "https://agentcoredatasets685197708687-284077270265-us-west-2-an.s3.us-west-2.amazonaws.com/685197708687/datasets/agentcore_cli_dataset_fixture-qhP7ipAkaf/draft/dataset.jsonl",
18+
"downloadUrlExpiresAt": {
19+
"$date": "2026-08-13T14:27:05.644Z"
20+
}
21+
}

‎src/handlers/eval/dataset/__fixtures__/ListDatasetsCommand.23f97c9dcdd6350b.json‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
{
22
"datasets": [
33
{
4-
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-jzVpQaA5It",
5-
"datasetId": "agentcore_cli_dataset_fixture-jzVpQaA5It",
4+
"datasetArn": "arn:aws:bedrock-agentcore:us-west-2:685197708687:dataset/agentcore_cli_dataset_fixture-qhP7ipAkaf",
5+
"datasetId": "agentcore_cli_dataset_fixture-qhP7ipAkaf",
66
"datasetName": "agentcore_cli_dataset_fixture",
77
"status": "ACTIVE",
88
"schemaType": "AGENTCORE_EVALUATION_PREDEFINED_V1",
99
"exampleCount": 2,
1010
"createdAt": {
11-
"$date": "2026-08-07T19:02:41.330Z"
11+
"$date": "2026-08-13T14:21:59.997Z"
1212
},
1313
"updatedAt": {
14-
"$date": "2026-08-07T19:02:58.844Z"
14+
"$date": "2026-08-13T14:22:17.603Z"
1515
},
1616
"description": "Recorded fixture for the dataset commands",
1717
"draftStatus": "MODIFIED"

0 commit comments

Comments
 (0)