Skip to content

Commit 4d7834a

Browse files
jesseturner21jesseturner21
authored
Merge pull request #1287 from aws/fix/harness-model-maxTokens-invoke-override
fix(invoke): preserve model maxTokens when overriding model at harness invoke
2 parents aa8e049 + 36a88c1 commit 4d7834a

4 files changed

Lines changed: 216 additions & 16 deletions

File tree

‎src/cli/aws/agentcore-harness.ts‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,34 @@ import { randomUUID } from 'node:crypto';
1717

1818
export type HarnessStatus = 'CREATING' | 'READY' | 'UPDATING' | 'DELETING' | 'DELETED' | 'FAILED';
1919

20+
export interface BedrockModelConfig {
21+
modelId: string;
22+
temperature?: number;
23+
topP?: number;
24+
maxTokens?: number;
25+
}
26+
27+
export interface OpenAiModelConfig {
28+
modelId: string;
29+
apiKeyArn?: string;
30+
temperature?: number;
31+
topP?: number;
32+
maxTokens?: number;
33+
}
34+
35+
export interface GeminiModelConfig {
36+
modelId: string;
37+
apiKeyArn?: string;
38+
temperature?: number;
39+
topP?: number;
40+
topK?: number;
41+
maxTokens?: number;
42+
}
43+
2044
export interface HarnessModelConfiguration {
21-
bedrockModelConfig?: { modelId: string };
22-
openAiModelConfig?: { modelId: string; apiKeyArn?: string };
23-
geminiModelConfig?: { modelId: string; apiKeyArn?: string };
45+
bedrockModelConfig?: BedrockModelConfig;
46+
openAiModelConfig?: OpenAiModelConfig;
47+
geminiModelConfig?: GeminiModelConfig;
2448
}
2549

2650
export type HarnessSystemPrompt = { text: string }[];

‎src/cli/commands/invoke/__tests__/build-harness-base-opts.test.ts‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
import { buildHarnessBaseOpts } from '../action.js';
2+
import type { InvokeOptions } from '../types.js';
3+
import { describe, expect, it } from 'vitest';
4+
5+
describe('buildHarnessBaseOpts', () => {
6+
describe('preserves model inference params from harness spec when overriding model', () => {
7+
it('bedrock: includes temperature, topP, and maxTokens', () => {
8+
const options: InvokeOptions = { modelId: 'anthropic.claude-v3' };
9+
const harnessSpec = {
10+
provider: 'bedrock' as const,
11+
modelId: 'anthropic.claude-v2',
12+
temperature: 0.7,
13+
topP: 0.9,
14+
maxTokens: 500,
15+
};
16+
17+
const result = buildHarnessBaseOpts(options, harnessSpec);
18+
19+
expect(result.model).toEqual({
20+
bedrockModelConfig: {
21+
modelId: 'anthropic.claude-v3',
22+
temperature: 0.7,
23+
topP: 0.9,
24+
maxTokens: 500,
25+
},
26+
});
27+
});
28+
29+
it('open_ai: includes temperature, topP, maxTokens, and apiKeyArn', () => {
30+
const options: InvokeOptions = { modelId: 'gpt-5' };
31+
const harnessSpec = {
32+
provider: 'open_ai' as const,
33+
modelId: 'gpt-4',
34+
apiKeyArn: 'arn:aws:secretsmanager:us-east-1:123:secret:key',
35+
temperature: 0.5,
36+
topP: 0.8,
37+
maxTokens: 2048,
38+
};
39+
40+
const result = buildHarnessBaseOpts(options, harnessSpec);
41+
42+
expect(result.model).toEqual({
43+
openAiModelConfig: {
44+
modelId: 'gpt-5',
45+
apiKeyArn: 'arn:aws:secretsmanager:us-east-1:123:secret:key',
46+
temperature: 0.5,
47+
topP: 0.8,
48+
maxTokens: 2048,
49+
},
50+
});
51+
});
52+
53+
it('gemini: includes temperature, topP, topK, maxTokens, and apiKeyArn', () => {
54+
const options: InvokeOptions = { modelId: 'gemini-2.5-pro' };
55+
const harnessSpec = {
56+
provider: 'gemini' as const,
57+
modelId: 'gemini-2.5-flash',
58+
apiKeyArn: 'arn:aws:secretsmanager:us-east-1:123:secret:gemini',
59+
temperature: 0.3,
60+
topP: 0.95,
61+
topK: 0.5,
62+
maxTokens: 1024,
63+
};
64+
65+
const result = buildHarnessBaseOpts(options, harnessSpec);
66+
67+
expect(result.model).toEqual({
68+
geminiModelConfig: {
69+
modelId: 'gemini-2.5-pro',
70+
apiKeyArn: 'arn:aws:secretsmanager:us-east-1:123:secret:gemini',
71+
temperature: 0.3,
72+
topP: 0.95,
73+
topK: 0.5,
74+
maxTokens: 1024,
75+
},
76+
});
77+
});
78+
});
79+
80+
describe('omits undefined inference params', () => {
81+
it('bedrock: only includes modelId when no inference params set', () => {
82+
const options: InvokeOptions = { modelId: 'anthropic.claude-v3' };
83+
const harnessSpec = { provider: 'bedrock' as const, modelId: 'anthropic.claude-v2' };
84+
85+
const result = buildHarnessBaseOpts(options, harnessSpec);
86+
87+
expect(result.model).toEqual({
88+
bedrockModelConfig: { modelId: 'anthropic.claude-v3' },
89+
});
90+
});
91+
92+
it('open_ai: omits apiKeyArn and inference params when not set', () => {
93+
const options: InvokeOptions = { modelId: 'gpt-5', modelProvider: 'open_ai' };
94+
const harnessSpec = { provider: 'open_ai' as const, modelId: 'gpt-4' };
95+
96+
const result = buildHarnessBaseOpts(options, harnessSpec);
97+
98+
expect(result.model).toEqual({
99+
openAiModelConfig: { modelId: 'gpt-5' },
100+
});
101+
});
102+
});
103+
104+
describe('CLI options take precedence for apiKeyArn', () => {
105+
it('uses CLI apiKeyArn over harness spec', () => {
106+
const options: InvokeOptions = {
107+
modelId: 'gpt-5',
108+
apiKeyArn: 'arn:aws:secretsmanager:us-east-1:123:secret:cli-key',
109+
};
110+
const harnessSpec = {
111+
provider: 'open_ai' as const,
112+
modelId: 'gpt-4',
113+
apiKeyArn: 'arn:aws:secretsmanager:us-east-1:123:secret:spec-key',
114+
maxTokens: 1000,
115+
};
116+
117+
const result = buildHarnessBaseOpts(options, harnessSpec);
118+
119+
expect(result.model!.openAiModelConfig!.apiKeyArn).toBe('arn:aws:secretsmanager:us-east-1:123:secret:cli-key');
120+
expect(result.model!.openAiModelConfig!.maxTokens).toBe(1000);
121+
});
122+
});
123+
124+
describe('does not set model when no model override options provided', () => {
125+
it('returns empty opts when no model-related options are set', () => {
126+
const options: InvokeOptions = {};
127+
const harnessSpec = {
128+
provider: 'bedrock' as const,
129+
modelId: 'anthropic.claude-v2',
130+
maxTokens: 500,
131+
};
132+
133+
const result = buildHarnessBaseOpts(options, harnessSpec);
134+
135+
expect(result.model).toBeUndefined();
136+
});
137+
});
138+
139+
describe('harness-level execution limits', () => {
140+
it('forwards maxTokens, maxIterations, and timeoutSeconds from CLI options', () => {
141+
const options: InvokeOptions = {
142+
maxTokens: 100,
143+
maxIterations: 10,
144+
harnessTimeout: 30,
145+
};
146+
147+
const result = buildHarnessBaseOpts(options);
148+
149+
expect(result.maxTokens).toBe(100);
150+
expect(result.maxIterations).toBe(10);
151+
expect(result.timeoutSeconds).toBe(30);
152+
});
153+
});
154+
});

‎src/cli/commands/invoke/action.ts‎

Lines changed: 34 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { ConfigIO } from '../../../lib';
2-
import type { AgentCoreProjectSpec, AwsDeploymentTargets, DeployedState } from '../../../schema';
2+
import type { AgentCoreProjectSpec, AwsDeploymentTargets, DeployedState, HarnessModel } from '../../../schema';
33
import {
44
buildAguiRunInput,
55
executeBashCommand,
@@ -512,30 +512,52 @@ export async function handleInvoke(context: InvokeContext, options: InvokeOption
512512
// Shared Harness Helpers
513513
// ============================================================================
514514

515-
interface HarnessModel {
516-
provider?: string;
517-
modelId?: string;
518-
apiKeyArn?: string;
519-
}
520-
521-
function buildHarnessBaseOpts(
515+
export function buildHarnessBaseOpts(
522516
options: InvokeOptions,
523-
harnessSpec?: HarnessModel
517+
harnessSpec?: Partial<HarnessModel>
524518
): Partial<import('../../aws/agentcore-harness').InvokeHarnessOptions> {
525519
const baseOpts: Partial<import('../../aws/agentcore-harness').InvokeHarnessOptions> = {};
526520
if (options.modelId || options.modelProvider || options.apiKeyArn) {
527521
const provider = options.modelProvider ?? harnessSpec?.provider;
528522
const modelId = options.modelId ?? harnessSpec?.modelId ?? '';
529523
const apiKeyArn = options.apiKeyArn ?? harnessSpec?.apiKeyArn;
524+
const temperature = harnessSpec?.temperature;
525+
const topP = harnessSpec?.topP;
526+
const topK = harnessSpec?.topK;
527+
const modelMaxTokens = harnessSpec?.maxTokens;
530528
switch (provider) {
531529
case 'open_ai':
532-
baseOpts.model = { openAiModelConfig: { modelId, ...(apiKeyArn && { apiKeyArn }) } };
530+
baseOpts.model = {
531+
openAiModelConfig: {
532+
modelId,
533+
...(apiKeyArn && { apiKeyArn }),
534+
...(temperature !== undefined && { temperature }),
535+
...(topP !== undefined && { topP }),
536+
...(modelMaxTokens !== undefined && { maxTokens: modelMaxTokens }),
537+
},
538+
};
533539
break;
534540
case 'gemini':
535-
baseOpts.model = { geminiModelConfig: { modelId, ...(apiKeyArn && { apiKeyArn }) } };
541+
baseOpts.model = {
542+
geminiModelConfig: {
543+
modelId,
544+
...(apiKeyArn && { apiKeyArn }),
545+
...(temperature !== undefined && { temperature }),
546+
...(topP !== undefined && { topP }),
547+
...(topK !== undefined && { topK }),
548+
...(modelMaxTokens !== undefined && { maxTokens: modelMaxTokens }),
549+
},
550+
};
536551
break;
537552
default:
538-
baseOpts.model = { bedrockModelConfig: { modelId } };
553+
baseOpts.model = {
554+
bedrockModelConfig: {
555+
modelId,
556+
...(temperature !== undefined && { temperature }),
557+
...(topP !== undefined && { topP }),
558+
...(modelMaxTokens !== undefined && { maxTokens: modelMaxTokens }),
559+
},
560+
};
539561
break;
540562
}
541563
}

‎src/schema/schemas/agentcore-project.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ export type { ABTestMode, TargetRef, GatewayFilter, PerVariantOnlineEvaluationCo
7070
export { ABTestModeSchema, TargetRefSchema, GatewayFilterSchema } from './primitives/ab-test';
7171
export type { HttpGatewayTarget } from './primitives/http-gateway';
7272
export { HttpGatewayTargetSchema } from './primitives/http-gateway';
73-
export type { HarnessGatewayOutboundAuth, HarnessSpec, HarnessModelProvider } from './primitives/harness';
73+
export type { HarnessGatewayOutboundAuth, HarnessModel, HarnessSpec, HarnessModelProvider } from './primitives/harness';
7474
export {
7575
GatewayOAuthGrantTypeSchema,
7676
HarnessGatewayOutboundAuthSchema,

0 commit comments

Comments
 (0)