Skip to content

Commit 2ff0b0f

Browse files
AshishKumar4AshishKumar4
committed
refactor: restructure inference context to separate metadata from runtime state
- Split InferenceContext into InferenceMetadata (persistent) and InferenceContext (runtime) - Move userModelConfigs and runtimeOverrides from state to behavior instance properties - Add credentialsToRuntimeOverrides utility for SDK credential payload conversion - Update all agent behaviors to use getInferenceContext() for runtime context assembly - Migrate legacy inferenceContext field to metadata in state migration
1 parent 516f92c commit 2ff0b0f

File tree

16 files changed

+266
-90
lines changed

16 files changed

+266
-90
lines changed

‎worker/agents/assistants/realtimeCodeFixer.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -490,7 +490,7 @@ ${block.error}
490490

491491
const llmResponse = await infer({
492492
env: this.env,
493-
metadata: this.inferenceContext,
493+
metadata: this.inferenceContext.metadata,
494494
modelName: AGENT_CONFIG['realtimeCodeFixer'].name,
495495
reasoning_effort: 'low',
496496
temperature: 0.0,

‎worker/agents/core/behaviors/agentic.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ export class AgenticCodingBehavior extends BaseCodingBehavior<AgenticState> impl
8888
lastPackageJson: packageJson,
8989
sessionId: sandboxSessionId!,
9090
hostname,
91-
inferenceContext,
91+
metadata: inferenceContext.metadata,
9292
projectType: this.projectType,
9393
behaviorType: 'agentic'
9494
});

‎worker/agents/core/behaviors/base.ts‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import { BaseSandboxService } from '../../../services/sandbox/BaseSandboxService
1818
import { getTemplateImportantFiles } from '../../../services/sandbox/utils';
1919
import { createScratchTemplateDetails } from '../../utils/templates';
2020
import { WebSocketMessageData, WebSocketMessageType } from '../../../api/websocketTypes';
21-
import { InferenceContext } from '../../inferutils/config.types';
21+
import { AgentActionKey, InferenceContext, InferenceRuntimeOverrides, ModelConfig } from '../../inferutils/config.types';
2222
import { ModelConfigService } from '../../../database/services/ModelConfigService';
2323
import { fixProjectIssues } from '../../../services/code-fixer';
2424
import { FastCodeFixerOperation } from '../../operations/PostPhaseCodeFixer';
@@ -68,6 +68,9 @@ export abstract class BaseCodingBehavior<TState extends BaseProjectState>
6868
protected deepDebugConversationId: string | null = null;
6969

7070
protected staticAnalysisCache: StaticAnalysisResponse | null = null;
71+
72+
protected userModelConfigs?: Record<AgentActionKey, ModelConfig>;
73+
protected runtimeOverrides?: InferenceRuntimeOverrides;
7174

7275
protected operations: BaseCodingOperations = {
7376
regenerateFile: new FileRegenerationOperation(),
@@ -229,7 +232,7 @@ export abstract class BaseCodingBehavior<TState extends BaseProjectState>
229232
query: this.state.query,
230233
blueprint: this.state.blueprint,
231234
template: this.getTemplateDetails(),
232-
inferenceContext: this.state.inferenceContext
235+
inferenceContext: this.getInferenceContext()
233236
});
234237
}
235238
return this.projectSetupAssistant;
@@ -247,6 +250,22 @@ export abstract class BaseCodingBehavior<TState extends BaseProjectState>
247250
return this.generationPromise !== null;
248251
}
249252

253+
getUserModelConfigs(): Record<AgentActionKey, ModelConfig> | undefined {
254+
return this.userModelConfigs;
255+
}
256+
257+
setUserModelConfigs(configs: Record<AgentActionKey, ModelConfig> | undefined): void {
258+
this.userModelConfigs = configs;
259+
}
260+
261+
getRuntimeOverrides(): InferenceRuntimeOverrides | undefined {
262+
return this.runtimeOverrides;
263+
}
264+
265+
setRuntimeOverrides(overrides: InferenceRuntimeOverrides | undefined): void {
266+
this.runtimeOverrides = overrides;
267+
}
268+
250269
abstract getOperationOptions(): OperationOptions;
251270

252271
/**
@@ -293,8 +312,12 @@ export abstract class BaseCodingBehavior<TState extends BaseProjectState>
293312
const controller = this.getOrCreateAbortController();
294313

295314
return {
296-
...this.state.inferenceContext,
315+
metadata: this.state.metadata,
316+
enableFastSmartCodeFix: false, // TODO: Do we want to enable it via some config?
317+
enableRealtimeCodeFix: false, // TODO: Do we want to enable it via some config?
297318
abortSignal: controller.signal,
319+
userModelConfigs: this.getUserModelConfigs(),
320+
runtimeOverrides: this.getRuntimeOverrides(),
298321
};
299322
}
300323

@@ -496,7 +519,7 @@ export abstract class BaseCodingBehavior<TState extends BaseProjectState>
496519

497520
getModelConfigsInfo() {
498521
const modelService = new ModelConfigService(this.env);
499-
return modelService.getModelConfigsInfo(this.state.inferenceContext.userId);
522+
return modelService.getModelConfigsInfo(this.state.metadata.userId);
500523
}
501524

502525
getTotalFiles(): number {

‎worker/agents/core/behaviors/phasic.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ export class PhasicCodingBehavior extends BaseCodingBehavior<PhasicState> implem
114114
lastPackageJson: packageJson,
115115
sessionId: sandboxSessionId!,
116116
hostname,
117-
inferenceContext,
117+
metadata: inferenceContext.metadata,
118118
projectType: this.projectType,
119119
behaviorType: 'phasic'
120120
};
@@ -159,7 +159,7 @@ export class PhasicCodingBehavior extends BaseCodingBehavior<PhasicState> implem
159159
}
160160

161161
migrateStateIfNeeded(): void {
162-
const migratedState = StateMigration.migrateIfNeeded(this.state, this.logger) as PhasicState | null;
162+
const migratedState = StateMigration.migratePhasic(this.state, this.logger) as PhasicState | null;
163163
if (migratedState) {
164164
this.setState(migratedState);
165165
}
@@ -587,7 +587,7 @@ export class PhasicCodingBehavior extends BaseCodingBehavior<PhasicState> implem
587587
});
588588
},
589589
userContext,
590-
shouldAutoFix: this.state.inferenceContext.enableRealtimeCodeFix,
590+
shouldAutoFix: this.getInferenceContext().enableRealtimeCodeFix,
591591
fileChunkGeneratedCallback: streamChunks ? (filePath: string, chunk: string, format: 'full_content' | 'unified_diff') => {
592592
this.broadcast(WebSocketMessageResponses.FILE_CHUNK_GENERATED, {
593593
message: `Generating file: ${filePath}`,
@@ -627,7 +627,7 @@ export class PhasicCodingBehavior extends BaseCodingBehavior<PhasicState> implem
627627
? await runPreDeploySafetyGate({
628628
files: finalFiles,
629629
env: this.env,
630-
inferenceContext: this.state.inferenceContext,
630+
inferenceContext: this.getInferenceContext(),
631631
query: this.state.query,
632632
template: templateDetails,
633633
phase,
@@ -647,7 +647,7 @@ export class PhasicCodingBehavior extends BaseCodingBehavior<PhasicState> implem
647647
await this.deployToSandbox(safeFiles, false, phase.name, true);
648648
if (postPhaseFixing) {
649649
await this.applyDeterministicCodeFixes();
650-
if (this.state.inferenceContext.enableFastSmartCodeFix) {
650+
if (this.getInferenceContext().enableFastSmartCodeFix) {
651651
await this.applyFastSmartCodeFixes();
652652
}
653653
}

‎worker/agents/core/codingAgent.ts‎

Lines changed: 22 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { AgenticState, AgentState, BaseProjectState, CurrentDevState, MAX_PHASES
44
import { Blueprint } from "../schemas";
55
import { BaseCodingBehavior } from "./behaviors/base";
66
import { createObjectLogger, StructuredLogger } from '../../logger';
7-
import { InferenceContext } from "../inferutils/config.types";
7+
import { InferenceMetadata } from "../inferutils/config.types";
88
import { getMimeType } from 'hono/utils/mime';
99
import { normalizePath, isPathSafe } from '../../utils/pathUtils';
1010
import { FileManager } from '../services/implementations/FileManager';
@@ -28,6 +28,7 @@ import { RateLimitExceededError } from "shared/types/errors";
2828
import { ProjectObjective } from "./objectives/base";
2929
import { FileOutputType } from "../schemas";
3030
import { SecretsClient, type UserSecretsStoreStub } from '../../services/secrets/SecretsClient';
31+
import { StateMigration } from './stateMigration';
3132

3233
const DEFAULT_CONVERSATION_SESSION_ID = 'default';
3334

@@ -66,7 +67,7 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
6667
templateName: '',
6768
generatedFilesMap: {},
6869
conversationMessages: [],
69-
inferenceContext: {} as InferenceContext,
70+
metadata: {} as InferenceMetadata,
7071
shouldBeGenerating: false,
7172
sandboxInstanceId: undefined,
7273
commandsHistory: [],
@@ -123,7 +124,7 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
123124
): Promise<AgentState> {
124125
const { inferenceContext } = initArgs;
125126
const sandboxSessionId = DeploymentManager.generateNewSessionId();
126-
this.initLogger(inferenceContext.agentId, inferenceContext.userId, sandboxSessionId);
127+
this.initLogger(inferenceContext.metadata.agentId, inferenceContext.metadata.userId, sandboxSessionId);
127128

128129
// Infrastructure setup
129130
await this.gitInit();
@@ -148,6 +149,12 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
148149
* @param props - Optional props
149150
*/
150151
async onStart(props?: Record<string, unknown> | undefined): Promise<void> {
152+
// Run common migration FIRST, before any state access
153+
const migratedState = StateMigration.migrateCommon(this.state);
154+
if (migratedState) {
155+
this.setState(migratedState);
156+
}
157+
151158
this.logger().info(`Agent ${this.getAgentId()} session: ${this.state.sessionId} onStart`, { props });
152159

153160
this.logger().info('Bootstrapping CodeGeneratorAgent', { props });
@@ -171,7 +178,7 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
171178
this.logger().warn(`Agent ${this.getAgentId()} session: ${this.state.sessionId} onStart ignored, agent not initialized`);
172179
return;
173180
}
174-
181+
175182
// Ensure state is migrated for any previous versions
176183
this.behavior.migrateStateIfNeeded();
177184

@@ -191,14 +198,8 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
191198

192199
// Load the latest user configs
193200
const modelConfigService = new ModelConfigService(this.env);
194-
const userConfigsRecord = await modelConfigService.getUserModelConfigs(this.state.inferenceContext.userId);
195-
this.setState({
196-
...this.state,
197-
inferenceContext: {
198-
...this.state.inferenceContext,
199-
userModelConfigs: userConfigsRecord,
200-
},
201-
});
201+
const userConfigsRecord = await modelConfigService.getUserModelConfigs(this.state.metadata.userId);
202+
this.behavior.setUserModelConfigs(userConfigsRecord);
202203
this.logger().info(`Agent ${this.getAgentId()} session: ${this.state.sessionId} onStart: User configs loaded successfully`, {userConfigsRecord});
203204
}
204205

@@ -240,13 +241,13 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
240241

241242
logger(): StructuredLogger {
242243
if (!this._logger) {
243-
this._logger = this.initLogger(this.getAgentId(), this.state.inferenceContext.userId, this.state.sessionId);
244+
this._logger = this.initLogger(this.getAgentId(), this.state.metadata.userId, this.state.sessionId);
244245
}
245246
return this._logger;
246247
}
247248

248249
getAgentId() {
249-
return this.state.inferenceContext.agentId;
250+
return this.state.metadata.agentId;
250251
}
251252

252253
getWebSockets(): WebSocket[] {
@@ -266,7 +267,7 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
266267

267268
private getSecretsClient(): SecretsClient {
268269
if (!this.secretsClient) {
269-
const userId = this.state.inferenceContext.userId;
270+
const userId = this.state.metadata.userId;
270271
const stub = this.env.UserSecretsStore.get(
271272
this.env.UserSecretsStore.idFromName(userId)
272273
) as unknown as UserSecretsStoreStub;
@@ -347,8 +348,8 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
347348
// Save the app to database (authenticated users only)
348349
const appService = new AppService(this.env);
349350
await appService.createApp({
350-
id: this.state.inferenceContext.agentId,
351-
userId: this.state.inferenceContext.userId,
351+
id: this.state.metadata.agentId,
352+
userId: this.state.metadata.userId,
352353
sessionToken: null,
353354
title: this.state.blueprint.title || this.state.query.substring(0, 100),
354355
description: this.state.blueprint.description,
@@ -360,12 +361,12 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
360361
createdAt: new Date(),
361362
updatedAt: new Date()
362363
});
363-
this.logger().info(`App saved successfully to database for agent ${this.state.inferenceContext.agentId}`, {
364-
agentId: this.state.inferenceContext.agentId,
365-
userId: this.state.inferenceContext.userId,
364+
this.logger().info(`App saved successfully to database for agent ${this.state.metadata.agentId}`, {
365+
agentId: this.state.metadata.agentId,
366+
userId: this.state.metadata.userId,
366367
visibility: 'private'
367368
});
368-
this.logger().info(`Agent initialized successfully for agent ${this.state.inferenceContext.agentId}`);
369+
this.logger().info(`Agent initialized successfully for agent ${this.state.metadata.agentId}`);
369370
}
370371

371372
// ==========================================

‎worker/agents/core/state.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type { PhasicBlueprint, AgenticBlueprint, PhaseConceptType ,
22
FileOutputType,
33
Blueprint,
44
} from '../schemas';
5-
import type { InferenceContext } from '../inferutils/config.types';
5+
import type { InferenceMetadata } from '../inferutils/config.types';
66
import { BehaviorType, Plan, ProjectType } from './types';
77

88
export interface FileState extends FileOutputType {
@@ -45,7 +45,7 @@ export interface BaseProjectState {
4545
templateName: string | 'custom';
4646

4747
// Inference context
48-
inferenceContext: InferenceContext;
48+
readonly metadata: InferenceMetadata;
4949

5050
// Generation control
5151
shouldBeGenerating: boolean;

0 commit comments

Comments
 (0)