diff --git a/.changeset/gateway-tfy-metadata.md b/.changeset/gateway-tfy-metadata.md new file mode 100644 index 000000000..210d0599a --- /dev/null +++ b/.changeset/gateway-tfy-metadata.md @@ -0,0 +1,5 @@ +--- +"@truefoundry/trueforge": patch +--- + +Forward session, turn, and agent context as `x-tfy-metadata` on TrueFoundry-mode model and MCP gateway calls. diff --git a/packages/trueforge/src/apis/turns.ts b/packages/trueforge/src/apis/turns.ts index 07d6e5c95..b08613086 100644 --- a/packages/trueforge/src/apis/turns.ts +++ b/packages/trueforge/src/apis/turns.ts @@ -32,7 +32,7 @@ import { streamSSE } from 'hono/streaming'; import type { Logger } from 'winston'; import type { Authorizer } from '../auth/authorizer'; import type { ResolveRequestContext } from '../auth/identity'; -import configuration from '../config'; +import configuration, { isTrueFoundryModeEnabled } from '../config'; import type { AgentRecord, IAgentStore } from '../db/agentStore'; import type { IMcpServerWithAuthStore } from '../db/mcpServerStore'; import type { IModelProviderStore } from '../db/modelProviderStore'; @@ -51,10 +51,13 @@ import { StreamGoneError, type EventSubscription, type EventSubscriptionRegistry import { mintPeeredTurnId } from '../runtime/peeringIds'; import { validateSandboxFilePath } from '../runtime/sandboxFilePath'; import { + buildGatewayMetadata, buildTurnSandbox, + gatewayMetadataHeaders, getMcpConnection, getModelDetails, resolveSandboxProvider, + withGatewayMetadataHeaders, } from '../runtime/sessionResources'; import { checkSnapshotStatus } from '../sandbox/providerUtils'; import { canReadAgentBoundResource } from './agentAccess'; @@ -142,9 +145,9 @@ function createTurnResolver(deps: { modelProviderStore: IModelProviderStore; logger: Logger; signal: AbortSignal; - tenant_id: string; userRef: string; - sessionId: string; + session: SessionHandle; + turnId: string; }): TurnResourceResolver { const { mcpServerStore, @@ -154,10 +157,16 @@ function createTurnResolver(deps: { modelProviderStore, logger, signal, - tenant_id, userRef, - sessionId, + session, + turnId, } = deps; + const tenant_id = session.tenant_id; + const sessionId = session.session_id; + const metadataHeaders = isTrueFoundryModeEnabled() + ? gatewayMetadataHeaders(buildGatewayMetadata({ session, turnId })) + : {}; + return new TurnResourceResolver({ llm: async name => { const resolved = await getModelDetails({ @@ -167,7 +176,10 @@ function createTurnResolver(deps: { }); return { modelClient: new VercelAILLM({ - providerConfig: resolved.providerConfig, + providerConfig: { + ...resolved.providerConfig, + headers: { ...resolved.providerConfig.headers, ...metadataHeaders }, + }, logger, signal, }), @@ -187,7 +199,13 @@ function createTurnResolver(deps: { message: `Unknown MCP server "${name}" — not configured`, }); } - return connection; + return { + url: connection.url, + headers: withGatewayMetadataHeaders({ + headers: connection.headers, + metadataHeaders, + }), + }; }, mcpRequestTimeoutMs: configuration.MCP_REQUEST_TIMEOUT_MS, mcpConnectTimeoutMs: configuration.MCP_CONNECT_TIMEOUT_MS, @@ -373,6 +391,7 @@ export async function beginTurnExecution(params: { }): Promise<{ turn: TurnHandle; drainInput: TurnEventDrainInput }> { const { session, input, previous_turn_id: previousTurnId, userRef, deps } = params; const sessionId = session.session_id; + const turnId = mintPeeredTurnId(configuration.EXECUTOR_ID); const abortController = new AbortController(); const tenant_id = session.tenant_id; @@ -384,9 +403,9 @@ export async function beginTurnExecution(params: { modelProviderStore: deps.modelProviderStore, logger: deps.logger, signal: abortController.signal, - tenant_id, userRef, - sessionId, + session, + turnId, }); // First turn only: derive the title from the first user message. The store @@ -394,7 +413,7 @@ export async function beginTurnExecution(params: { const title = session.record.last_turn_id ? undefined : deriveSessionTitle(input); const turn = await session.createTurn({ - turn_id: mintPeeredTurnId(configuration.EXECUTOR_ID), + turn_id: turnId, input, previous_turn_id: previousTurnId, signal: abortController.signal, diff --git a/packages/trueforge/src/runtime/sessionResources.ts b/packages/trueforge/src/runtime/sessionResources.ts index 9cd1c5450..57ef96ea2 100644 --- a/packages/trueforge/src/runtime/sessionResources.ts +++ b/packages/trueforge/src/runtime/sessionResources.ts @@ -1,4 +1,4 @@ -import type { AgentSpec } from '@truefoundry/trueforge-core/agent-session'; +import type { AgentSpec, SessionHandle } from '@truefoundry/trueforge-core/agent-session'; import { Sandbox, SkillMounter, @@ -28,6 +28,59 @@ export interface McpConnection { headers: RemoteMcpHeaders; } +/** Gateway header carrying stringified JSON metadata. */ +export const X_TFY_METADATA = 'x-tfy-metadata'; + +/** Prefix for harness-owned keys */ +export const TFG_METADATA_PREFIX = 'tfg'; + +export function buildGatewayMetadata(input: { session: SessionHandle; turnId: string }): Record { + // Session.metadata is intentionally omitted for now (Unicode-in-header risk); re-add later. + const metadata: Record = { + [`${TFG_METADATA_PREFIX}.session_id`]: input.session.session_id, + [`${TFG_METADATA_PREFIX}.turn_id`]: input.turnId, + }; + const { agent } = input.session; + if (agent.type === 'reference') { + metadata[`${TFG_METADATA_PREFIX}.agent_id`] = agent.id; + if (agent.name !== null) { + metadata[`${TFG_METADATA_PREFIX}.agent_name`] = agent.name; + } + } + return metadata; +} + +export function gatewayMetadataHeaders(metadata: Record): Record { + if (Object.keys(metadata).length === 0) { + return {}; + } + return { [X_TFY_METADATA]: JSON.stringify(metadata) }; +} + +/** + * Merge gateway metadata into MCP invoke headers. Preserves authRequired; + * metadata is applied after auth/per-server headers. + */ +export function withGatewayMetadataHeaders(input: { + headers: RemoteMcpHeaders; + metadataHeaders: Record; +}): RemoteMcpHeaders { + const { headers, metadataHeaders } = input; + if (Object.keys(metadataHeaders).length === 0) { + return headers; + } + if (typeof headers !== 'function') { + return { ...headers, ...metadataHeaders }; + } + return async () => { + const result = await headers(); + if ('authRequired' in result) { + return result; + } + return { headers: { ...result.headers, ...metadataHeaders } }; + }; +} + /** Split `provider/model` FQN. Returns undefined when the shape is not exactly one slash. */ export function parseModelFqn(name: string): { providerName: string; modelName: string } | undefined { const slash = name.indexOf('/'); diff --git a/packages/trueforge/tests/unit/runtime/sessionResources.test.ts b/packages/trueforge/tests/unit/runtime/sessionResources.test.ts index 7c2997a0d..9bf1c2566 100644 --- a/packages/trueforge/tests/unit/runtime/sessionResources.test.ts +++ b/packages/trueforge/tests/unit/runtime/sessionResources.test.ts @@ -1,4 +1,10 @@ -import { AgentSpecSchema } from '@truefoundry/trueforge-core/agent-session'; +import { + AgentSpecSchema, + InMemorySessionStore, + Sessions, + type SessionAgent, + type SessionHandle, +} from '@truefoundry/trueforge-core/agent-session'; import { HTTPException } from 'hono/http-exception'; import { validateGitAgentSkills } from '../../../src/db/gitSkillMounts'; import { migrateSqliteToLatest } from '../../../src/db/migrateSqlite'; @@ -8,10 +14,78 @@ import { SqliteMcpServerStore } from '../../../src/db/sqlite/mcp-server-store/Sq import { SqliteModelProviderStore } from '../../../src/db/sqlite/model-provider-store/SqliteModelProviderStore'; import { SqliteSandboxProviderStore } from '../../../src/db/sqlite/sandbox-provider-store/SqliteSandboxProviderStore'; import { SqliteSkillStore } from '../../../src/db/sqlite/skill-store/SqliteSkillStore'; -import { getModelDetails, localSandboxSessionSegment, validateAgentSpec } from '../../../src/runtime/sessionResources'; +import { + buildGatewayMetadata, + getModelDetails, + localSandboxSessionSegment, + TFG_METADATA_PREFIX, + validateAgentSpec, + withGatewayMetadataHeaders, + X_TFY_METADATA, +} from '../../../src/runtime/sessionResources'; import { setCachedLocalSandboxSupport } from '../../../src/sandbox/localRuntime'; import type { ReasoningEffort } from '../../../src/schemas/modelProvider'; +async function createGatewayMetadataSession(input: { agent: SessionAgent }): Promise { + const sessions = new Sessions({ sessionStore: new InMemorySessionStore() }); + return sessions.create({ + tenant_id: 'tenant-1', + session_id: 'sess-1', + created_by_subject: { subject_id: 'user-1', subject_type: 'user', subject_display_name: 'user-1' }, + agent: input.agent, + metadata: {}, + external_id: null, + }); +} + +describe('buildGatewayMetadata', () => { + it('stamps session/turn/agent fields only', async () => { + const session = await createGatewayMetadataSession({ + agent: { type: 'reference', id: 'agent-1', name: 'my-agent' }, + }); + + expect(buildGatewayMetadata({ session, turnId: 'turn-1' })).toEqual({ + [`${TFG_METADATA_PREFIX}.session_id`]: 'sess-1', + [`${TFG_METADATA_PREFIX}.turn_id`]: 'turn-1', + [`${TFG_METADATA_PREFIX}.agent_id`]: 'agent-1', + [`${TFG_METADATA_PREFIX}.agent_name`]: 'my-agent', + }); + }); +}); + +describe('withGatewayMetadataHeaders', () => { + it('merges into async header resolvers and preserves authRequired', async () => { + const withAuth = withGatewayMetadataHeaders({ + headers: async () => ({ headers: { Authorization: 'Bearer t' } }), + metadataHeaders: { [X_TFY_METADATA]: '{"k":"v"}' }, + }); + expect(typeof withAuth).toBe('function'); + if (typeof withAuth !== 'function') { + throw new Error('expected async header resolver'); + } + await expect(withAuth()).resolves.toEqual({ + headers: { + Authorization: 'Bearer t', + [X_TFY_METADATA]: '{"k":"v"}', + }, + }); + + const authRequired = withGatewayMetadataHeaders({ + headers: async () => ({ + authRequired: { servers: [{ id: 'mcp', name: 'mcp', auth_url: 'https://auth.example' }] }, + }), + metadataHeaders: { [X_TFY_METADATA]: '{"k":"v"}' }, + }); + expect(typeof authRequired).toBe('function'); + if (typeof authRequired !== 'function') { + throw new Error('expected async header resolver'); + } + await expect(authRequired()).resolves.toEqual({ + authRequired: { servers: [{ id: 'mcp', name: 'mcp', auth_url: 'https://auth.example' }] }, + }); + }); +}); + describe('localSandboxSessionSegment', () => { it('keeps a single-segment session id and rejects missing or unsafe values', () => { expect(localSandboxSessionSegment('sess_1')).toBe('sess_1');