@db3.ai/app 0.1.0-beta.1 → 0.1.0

package/src/ai/Agent.ts

@@ -1,3 +1,4 @@
+import { AIProviderDeferredError, AIProviderStoppedError } from './AIProviderAdmission';
 import { aiExecutionContext } from './AiExecutionContext';
 import { Agent as SdkAgent, Runner, generateTraceId, type AgentInputItem, type AgentOutputType, type FunctionCallItem, type FunctionCallResultItem, type ModelRequest, type ModelSettings, type RunStreamEvent, type Usage } from '@openai/agents';
 import { isQueueRetryLaterError } from '@db3.ai/app/queue';
@@ -12,9 +13,11 @@ import { AI_PROVIDER, AIQuotaDeferredError, AIRateLimitDeferredError, isFailover
 import { AIAllowanceExceededError } from './AIErrors.js';
 import type { AIUsageCharge } from './contracts/AI';
 import { estimateTokensFromText } from './Estimates.js';
-import { calculateAIHostedToolCost, calculateAIRequestCostUSD, calculateAIRequestEntriesCostUSD, combineAIRequestCostUSD, type AIRequestCostUsageEntry } from './Models.js';
+import type { AIServiceTier } from './contracts/AIServiceTier';
+import { assertAIServiceTier, effectiveAIServiceTier } from './aiServiceTier';
+import { calculateAIHostedToolCost, calculateAIRequestCostUSD, calculateAIRequestEntriesCostUSD, combineAIRequestCostUSD, type AIRequestCostUsageEntry } from './modelPricing.js';
 import { AgentRunJob } from './AgentRunJob.js';
-import { openAIQuotaRetryDecision, withOpenAIQuotaRetryMetadata, withoutOpenAIQuotaRetryMetadata, type OpenAIQuotaRetryDecision } from './OpenAIQuotaRetry.js';
+import { withoutOpenAIQuotaRetryMetadata } from './OpenAIQuotaRetry.js';
 import { agentCitationFromStreamRecord, agentCitationsFromFinalModelResponse, rememberAgentCitation } from './agentCitations.js';
 import { serializeQueuedAgent } from './registry.js';
 import { agentToolResultFromStructuredOutput } from './toolHelpers.js';
@@ -25,6 +28,8 @@ import type * as agent from './contracts/Agent';
  * Fully prepared run inputs and persistence state ready for SDK execution.
  */
 interface PreparedAgentRun {
+	/** Tier captured at admission and retained by queued runs. */
+	serviceTier: AIServiceTier;
 	/** User-visible task message supplied to the agent. */
 	message: string;
 
@@ -115,7 +120,8 @@ interface ResolvedQueueInput {
  * Shared persisted Agents SDK runtime for interactive and queued application agents.
  */
 export abstract class Agent<TContext extends agent.BaseAgentContext = agent.BaseAgentContext> implements Serializable<TContext> {
-	protected readonly maxTurns: number = 8;
+	/** Maximum model turns per run; null lets the agent finish without a turn ceiling. */
+	protected readonly maxTurns: number | null = 8;
 	protected readonly maxPreviousMessages: number = 50;
 	/** Provider reasoning effort used for every run of this agent. */
 	protected readonly reasoningEffort: agent.AgentReasoningEffort | null = null;
@@ -124,6 +130,8 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 	protected readonly workflowName: string = this.constructor.name;
 	/** Optional agent-specific model. Omit to use the global `app().ai.model`. */
 	protected readonly model?: string;
+	/** OpenAI processing tier for this agent. Override to opt into local Flex experiments. */
+	protected readonly serviceTier: AIServiceTier = 'default';
 	private currentInput: AgentRunInputPayload | null = null;
 
 	/**
@@ -366,31 +374,13 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 
 			return await this.executePreparedRun(prepared, async () => {});
 		} catch (error) {
-			if (isQueueRetryLaterError(error)) throw error;
-
-			const quotaRetry = openAIQuotaRetryDecision(error, prepared.state.aiRequest);
-
-			if (quotaRetry && !quotaRetry.expired) {
-				await this.deferTrackedRun(prepared.state, quotaRetry);
-
-				throw new AIQuotaDeferredError(
-					quotaRetry.delaySeconds,
-					quotaRetry.retryAt,
-					quotaRetry.retryUntil,
-					quotaRetry.deferralCount,
-					quotaRetry.message,
-				);
+			if (error instanceof AIProviderDeferredError) {
+				await this.deferTrackedRun(prepared.state, error);
+				throw error;
 			}
-
+			if (isQueueRetryLaterError(error)) throw error;
 			const providerError = openAIProviderError(error);
 
-			if (quotaRetry) {
-				prepared.state.aiRequest.metadata = withOpenAIQuotaRetryMetadata(
-					prepared.state.aiRequest.metadata,
-					quotaRetry.metadata,
-				);
-			}
-
 			await this.failTrackedRun(
 				prepared.state,
 				error instanceof Error ? error.message : `Unable to run ${this.agentName}.`,
@@ -698,7 +688,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 			input: invocation.input,
 			maxTurns: this.maxTurns,
 			model,
-			modelSettings: invocation.modelSettings,
+			modelSettings: modelSettingsForProvider(invocation.modelSettings, AI_PROVIDER.openai, this.serviceTier),
 		});
 	}
 
@@ -894,6 +884,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 		return {
 			message,
 			model,
+			serviceTier: this.serviceTier,
 			traceId,
 			previousMessages,
 			invocation,
@@ -923,6 +914,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 		return {
 			message: payload.message,
 			model: payload.model,
+			serviceTier: payload.serviceTier ?? 'default',
 			traceId: payload.traceId,
 			previousMessages,
 			invocation,
@@ -957,6 +949,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 		const instructions = invocation.instructions;
 		const modelInput = invocation.input;
 		const attempts = app().ai.resolveProviders(this.provider(), prepared.model);
+		for (const attempt of attempts) assertAIServiceTier(attempt.provider, prepared.serviceTier);
 		const attemptFailures: AIProviderAttemptFailure[] = [];
 
 		await this.attachPromptRequestPayload(state, {
@@ -977,13 +970,15 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 			const attemptStartMs = Date.now();
 			let rateLimitLease: AIRateLimitLease | null = null;
 			let sawStreamEvent = false;
+			let completedModelTurn = false;
 			let activeUsage: Usage | null = null;
 			let attemptState: agent.AgentPersistenceState | null = null;
 			let attemptSettled = false;
 			let providerStarted = false;
+			const responsePricing: Array<{ model?: string; serviceTier: string | null }> = [];
 
 			try {
-				attemptState = await this.startTrackedAttempt(state, attempt, attemptStartMs);
+				attemptState = await this.startTrackedAttempt(state, attempt, attemptStartMs, prepared.serviceTier);
 				rateLimitLease = await this.acquireProviderCapacity(
 					prepared,
 					instructions,
@@ -1001,8 +996,19 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 					providerStarted = true;
 					const runner = new Runner({
 						model: attempt.model,
-						modelProvider: app().ai.createAgentProvider(attempt, rateLimitLease),
-						modelSettings: modelSettingsForProvider(invocation.modelSettings, attempt.provider),
+						modelProvider: app().ai.createAgentProvider(attempt, rateLimitLease, prepared.serviceTier),
+						modelSettings: {
+							...modelSettingsForProvider(invocation.modelSettings, attempt.provider, prepared.serviceTier),
+							retry: {
+								maxRetries: 8,
+								/** Retains completed tools in this SDK run; never restarts the whole agent. */
+								policy: ({ error }) => {
+									if (!completedModelTurn) return false;
+									const delayMs = app().ai.providerAdmission.requestRetryDelay(attempt, agentRateLimitError(error));
+									return delayMs === null ? false : { retry: true, delayMs };
+								},
+							},
+						},
 						tracingDisabled: process.env.OPENAI_AGENTS_TRACING_DISABLED === 'true',
 						traceIncludeSensitiveData: false,
 						workflowName: this.workflowName,
@@ -1024,6 +1030,14 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 
 					for await (const event of stream) {
 						sawStreamEvent = true;
+						if (event.type === 'raw_model_stream_event' && event.data.type === 'response_done') {
+							completedModelTurn = true;
+							const providerData = recordValue(event.data.response.providerData);
+							responsePricing.push({
+								...(typeof providerData?.model === 'string' ? { model: providerData.model } : {}),
+								serviceTier: effectiveAIServiceTier(providerData?.service_tier, prepared.serviceTier),
+							});
+						}
 
 						const streamEvent = await this.streamEventFromSdkEvent(executionState, event);
 
@@ -1037,6 +1051,8 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 					}
 
 					await stream.completed;
+					const streamFailure = app().ai.providerAdmission?.requestFailure(attempt, null);
+					if (streamFailure) throw streamFailure;
 					const rawFinalOutput = this.finalOutputText(stream.finalOutput, assistantText);
 					const finalResponseCitations = agentCitationsFromFinalModelResponse(stream.rawResponses);
 					const citations = finalResponseCitations ?? normalizedAgentCitations(executionState.citations);
@@ -1054,6 +1070,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 						executionState,
 						activeUsage,
 						stream.lastResponseId ?? null,
+						responsePricing,
 					);
 					attemptSettled = true;
 
@@ -1102,13 +1119,18 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 					};
 				});
 			} catch (caught) {
-				const error = agentRateLimitError(caught);
+				const restored = agentRateLimitError(caught);
+				const failure = app().ai.providerAdmission?.requestFailure(attempt, restored) ?? restored;
+				// Restarting a turn after streamed output/tools can repeat completed work.
+				const error = failure instanceof AIProviderDeferredError && sawStreamEvent ? new AIProviderStoppedError('outage', failure.providerCode, { ...failure.diagnostics, stopStage: 'stream-output' }) : failure;
+				if (failure instanceof AIProviderDeferredError && error instanceof AIProviderStoppedError) error.providerStarted = failure.providerStarted;
 				if (attemptState && !attemptSettled) {
 					await this.failTrackedAttempt(
 						attemptState,
 						activeUsage,
 						providerStarted,
 						error,
+						responsePricing,
 					);
 				}
 
@@ -1151,12 +1173,14 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 	 * @param rootState - Logical run root and shared conversation state.
 	 * @param attempt - Provider and model selected for this execution.
 	 * @param startedAtMs - Attempt start timestamp in milliseconds.
+	 * @param serviceTier - Requested tier retained alongside this attempt.
 	 * @returns Persistence state scoped to the new provider attempt.
 	 */
 	private async startTrackedAttempt(
 		rootState: agent.AgentPersistenceState,
 		attempt: AIProviderAttempt,
 		startedAtMs: number,
+		serviceTier: AIServiceTier,
 	): Promise<agent.AgentPersistenceState> {
 		const aiRequest = await new (app().ai.models.request)({
 			conversation: rootState.conversation,
@@ -1168,6 +1192,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 			operation: 'agents.run.attempt',
 			status: AI_REQUEST_STATUS.pending,
 			request: {
+				serviceTier,
 				rootAiRequestId: rootState.rootAiRequest.id,
 				traceId: rootState.traceId,
 			},
@@ -1204,15 +1229,17 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 		state: agent.AgentPersistenceState,
 		sdkUsage: Usage,
 		lastResponseId: string | null,
+		responsePricing: Array<{ model?: string; serviceTier: string | null }>,
 	): Promise<void> {
 		const usage = agentRunUsage(sdkUsage);
-		const requestUsageEntries = agentRequestCostUsageEntries(sdkUsage);
+		const requestUsageEntries = agentRequestCostUsageEntries(sdkUsage, responsePricing, recordValue(state.aiRequest.request)?.serviceTier);
 		const cost = agentAttemptCost(
 			state.aiRequest.provider,
 			state.model,
 			usage,
 			requestUsageEntries,
 			state.hostedToolUsage,
+			recordValue(state.aiRequest.request)?.serviceTier,
 		);
 
 		state.aiRequest.assign({
@@ -1250,9 +1277,10 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 		sdkUsage: Usage | null,
 		providerStarted: boolean,
 		error: unknown,
+		responsePricing: Array<{ model?: string; serviceTier: string | null }>,
 	): Promise<void> {
 		const usage = sdkUsage ? agentRunUsage(sdkUsage) : emptyAgentRunUsage();
-		const requestUsageEntries = sdkUsage ? agentRequestCostUsageEntries(sdkUsage) : [];
+		const requestUsageEntries = sdkUsage ? agentRequestCostUsageEntries(sdkUsage, responsePricing, recordValue(state.aiRequest.request)?.serviceTier) : [];
 		const hasBillableUsage = sdkUsage
 			? hasBillableAgentAttemptUsage(sdkUsage, state.hostedToolUsage)
 			: Object.values(state.hostedToolUsage).some(count => Number.isFinite(count) && count > 0);
@@ -1263,6 +1291,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 				usage,
 				requestUsageEntries,
 				state.hostedToolUsage,
+				recordValue(state.aiRequest.request)?.serviceTier,
 			)
 			: emptyAgentAttemptCost(providerStarted);
 		const message = error instanceof Error ? error.message : `Unable to run ${this.agentName}.`;
@@ -1910,36 +1939,13 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 		};
 	}
 
-	/**
-	 * Keeps a quota-blocked queued run pending and records its next retry.
-	 *
-	 * @param state - Active run persistence state.
-	 * @param quotaRetry - Bounded exponential quota retry decision.
-	 */
-	private async deferTrackedRun(
-		state: agent.AgentPersistenceState,
-		quotaRetry: OpenAIQuotaRetryDecision,
-	): Promise<void> {
-		state.aiRequest.assign({
-			status: AI_REQUEST_STATUS.pending,
-			completedAt: null,
-			durationMs: null,
-			errorCode: quotaRetry.code,
-			errorMessage: quotaRetry.message,
-			metadata: withOpenAIQuotaRetryMetadata(
-				state.aiRequest.metadata,
-				quotaRetry.metadata,
-			),
-		});
-
-		await state.aiRequest.save();
-		await this.touchConversation(state, {
-			lastTraceId: state.traceId,
-			lastStatus: AI_REQUEST_STATUS.pending,
-			lastError: quotaRetry.message,
-			nextRetryAt: quotaRetry.retryAt.toISOString(),
-			retryUntil: quotaRetry.retryUntil.toISOString(),
+	/** Keeps unstarted queued work visibly pending on the shared account deadline. */
+	private async deferTrackedRun(state: agent.AgentPersistenceState, error: AIProviderDeferredError): Promise<void> {
+		state.rootAiRequest.assign({ status: AI_REQUEST_STATUS.pending, completedAt: null, errorCode: 'ai_provider_deferred', errorMessage: error.message,
+			metadata: { ...(state.rootAiRequest.metadata ?? {}), providerAdmission: { retryAt: error.retryAt.toISOString(), retryUntil: error.retryUntil.toISOString() } },
 		});
+		await state.rootAiRequest.save();
+		await this.touchConversation(state, { lastTraceId: state.traceId, lastStatus: AI_REQUEST_STATUS.pending, lastError: error.message, nextRetryAt: error.retryAt.toISOString(), retryUntil: error.retryUntil.toISOString() });
 	}
 
 	/**
@@ -2044,6 +2050,7 @@ export abstract class Agent<TContext extends agent.BaseAgentContext = agent.Base
 			aiRequestId: prepared.state.aiRequest.id ? String(prepared.state.aiRequest.id) : null,
 			traceId: prepared.traceId,
 			model: prepared.model,
+			serviceTier: prepared.serviceTier,
 		};
 	}
 
@@ -2108,7 +2115,7 @@ function agentRateLimitError(error: unknown): unknown {
 	const seen = new Set<unknown>();
 	let current = error;
 	while (current instanceof Error && !seen.has(current)) {
-		if (current instanceof AIRateLimitDeferredError) return current;
+		if (isQueueRetryLaterError(current) || current instanceof AIProviderStoppedError) return current;
 		seen.add(current);
 		current = current.cause;
 	}
@@ -2143,8 +2150,8 @@ interface OpenAIResponseRequestPreview {
 	/** Complete first-turn provider input. */
 	input: agent.AgentModelInput;
 
-	/** Maximum turns configured for the live agent run. */
-	maxTurns: number;
+	/** Maximum turns configured for the live run, or null for no turn ceiling. */
+	maxTurns: number | null;
 
 	/** OpenAI model name used for request defaults and the provider body. */
 	model: string;
@@ -3297,10 +3304,13 @@ function agentRunUsage(usage: Usage): agent.AgentRunUsage {
  * Normalizes the SDK's individual provider-request usage for accurate pricing.
  *
  * @param usage - Agents SDK aggregate usage object.
+ * @param responsePricing - Terminal response metadata in SDK request order.
+ * @param requestedTier - Saved selection used only for legacy Standard fallback.
  * @returns Individual request entries with cached-input details preserved.
  */
-function agentRequestCostUsageEntries(usage: Usage): AIRequestCostUsageEntry[] {
-	return (usage.requestUsageEntries ?? []).map(entry => ({
+function agentRequestCostUsageEntries(usage: Usage, responsePricing: Array<{ model?: string; serviceTier: string | null }>, requestedTier: unknown): AIRequestCostUsageEntry[] {
+	return (usage.requestUsageEntries ?? []).map((entry, index) => ({
+		...(responsePricing.length === usage.requestUsageEntries?.length ? responsePricing[index] : { serviceTier: effectiveAIServiceTier(undefined, requestedTier) }),
 		inputTokens: numberOrNull(entry.inputTokens),
 		outputTokens: numberOrNull(entry.outputTokens),
 		cachedTokens: tokenDetailValue(entry.inputTokensDetails, 'cached_tokens'),
@@ -3337,6 +3347,7 @@ interface AgentAttemptCost {
  * @param usage - Aggregate attempt usage.
  * @param requestUsageEntries - Per-provider-request usage entries.
  * @param hostedToolUsage - Completed hosted provider-tool calls.
+ * @param requestedTier - Saved selection used for single-request Standard fallback.
  * @returns Complete attempt cost and any unpriced dimensions.
  */
 function agentAttemptCost(
@@ -3345,6 +3356,7 @@ function agentAttemptCost(
 	usage: agent.AgentRunUsage,
 	requestUsageEntries: AIRequestCostUsageEntry[],
 	hostedToolUsage: Record<string, number>,
+	requestedTier: unknown,
 ): AgentAttemptCost {
 	const usesPerRequestUsage = usage.requests > 0
 		&& requestUsageEntries.length === usage.requests;
@@ -3354,7 +3366,7 @@ function agentAttemptCost(
 		? usesPerRequestUsage
 			? calculateAIRequestEntriesCostUSD(model, requestUsageEntries)
 			: usesSingleRequestAggregate
-				? calculateAIRequestCostUSD(model, usage)
+				? calculateAIRequestCostUSD(model, { ...usage, serviceTier: effectiveAIServiceTier(undefined, requestedTier) })
 				: null
 		: null;
 	const hostedToolCost = provider === AI_PROVIDER.openai
@@ -3470,20 +3482,22 @@ function emptyAgentRunUsage(): agent.AgentRunUsage {
 }
 
 /**
- * Applies Standard service tier only to direct OpenAI provider requests.
+ * Applies the explicitly selected service tier only to direct OpenAI requests.
  *
  * @param settings - Agent-owned model settings.
  * @param provider - Provider selected for the current attempt.
- * @returns Provider-safe settings with deterministic OpenAI billing tier.
+ * @param serviceTier - Explicit class selection, frozen in queued run payloads.
+ * @returns Provider-safe settings with the selected OpenAI tier.
  */
-function modelSettingsForProvider(settings: ModelSettings, provider: string): ModelSettings {
+function modelSettingsForProvider(settings: ModelSettings, provider: string, serviceTier: AIServiceTier): ModelSettings {
+	assertAIServiceTier(provider, serviceTier);
 	if (provider !== AI_PROVIDER.openai) return settings;
 
 	return {
 		...settings,
 		providerData: {
 			...(settings.providerData ?? {}),
-			service_tier: 'default',
+			service_tier: serviceTier,
 		},
 	};
 }