feat(antigravity): add thought signature support for multi-turn conversations

Gemini 3 Pro requires thoughtSignature on function call blocks in
subsequent requests. This commit:

- Add thought-signature-store for session-based signature storage
- Extract signature from both streaming (SSE) and non-streaming responses
- Inject signature into functionCall parts on subsequent requests
- Maintain consistent sessionId per fetch instance

Debug logging available via ANTIGRAVITY_DEBUG=1

🤖 GENERATED WITH ASSISTANCE OF [OhMyOpenCode](https://github.com/code-yeongyu/oh-my-opencode)
This commit is contained in:
YeonGyu-Kim
2025-12-13 01:16:53 +09:00
parent 3d273ff853
commit 5fd59afacf
5 changed files with 358 additions and 75 deletions

View File

@@ -21,9 +21,19 @@ import { ANTIGRAVITY_ENDPOINT_FALLBACKS } from "./constants"
import { fetchProjectContext, clearProjectContextCache } from "./project" import { fetchProjectContext, clearProjectContextCache } from "./project"
import { isTokenExpired, refreshAccessToken, parseStoredToken, formatTokenForStorage } from "./token" import { isTokenExpired, refreshAccessToken, parseStoredToken, formatTokenForStorage } from "./token"
import { transformRequest } from "./request" import { transformRequest } from "./request"
import { transformResponse, transformStreamingResponse, isStreamingResponse } from "./response" import {
transformResponse,
transformStreamingResponse,
isStreamingResponse,
extractSignatureFromSsePayload,
} from "./response"
import { normalizeToolsForGemini, type OpenAITool } from "./tools" import { normalizeToolsForGemini, type OpenAITool } from "./tools"
import { extractThinkingBlocks, shouldIncludeThinking, transformResponseThinking } from "./thinking" import { extractThinkingBlocks, shouldIncludeThinking, transformResponseThinking } from "./thinking"
import {
getThoughtSignature,
setThoughtSignature,
getOrCreateSessionId,
} from "./thought-signature-store"
import type { AntigravityTokens } from "./types" import type { AntigravityTokens } from "./types"
/** /**
@@ -65,14 +75,22 @@ function isRetryableError(status: number): boolean {
return false return false
} }
async function attemptFetch( interface AttemptFetchOptions {
endpoint: string, endpoint: string
url: string, url: string
init: RequestInit, init: RequestInit
accessToken: string, accessToken: string
projectId: string, projectId: string
sessionId: string
modelName?: string modelName?: string
thoughtSignature?: string
}
async function attemptFetch(
options: AttemptFetchOptions
): Promise<Response | null | "pass-through"> { ): Promise<Response | null | "pass-through"> {
const { endpoint, url, init, accessToken, projectId, sessionId, modelName, thoughtSignature } =
options
debugLog(`Trying endpoint: ${endpoint}`) debugLog(`Trying endpoint: ${endpoint}`)
try { try {
@@ -99,14 +117,16 @@ async function attemptFetch(
} }
} }
const transformed = transformRequest( const transformed = transformRequest({
url, url,
parsedBody, body: parsedBody,
accessToken, accessToken,
projectId, projectId,
sessionId,
modelName, modelName,
endpoint endpointOverride: endpoint,
) thoughtSignature,
})
const response = await fetch(transformed.url, { const response = await fetch(transformed.url, {
method: init.method || "POST", method: init.method || "POST",
@@ -129,16 +149,56 @@ async function attemptFetch(
} }
} }
/** interface GeminiResponsePart {
* Transform response with thinking extraction if applicable thoughtSignature?: string
*/ thought_signature?: string
functionCall?: Record<string, unknown>
text?: string
[key: string]: unknown
}
interface GeminiResponseCandidate {
content?: {
parts?: GeminiResponsePart[]
[key: string]: unknown
}
[key: string]: unknown
}
interface GeminiResponseBody {
candidates?: GeminiResponseCandidate[]
[key: string]: unknown
}
function extractSignatureFromResponse(parsed: GeminiResponseBody): string | undefined {
if (!parsed.candidates || !Array.isArray(parsed.candidates)) {
return undefined
}
for (const candidate of parsed.candidates) {
const parts = candidate.content?.parts
if (!parts || !Array.isArray(parts)) {
continue
}
for (const part of parts) {
const sig = part.thoughtSignature || part.thought_signature
if (sig && typeof sig === "string") {
return sig
}
}
}
return undefined
}
async function transformResponseWithThinking( async function transformResponseWithThinking(
response: Response, response: Response,
modelName: string modelName: string,
fetchInstanceId: string
): Promise<Response> { ): Promise<Response> {
const streaming = isStreamingResponse(response) const streaming = isStreamingResponse(response)
// Transform response based on streaming mode
let result let result
if (streaming) { if (streaming) {
result = await transformStreamingResponse(response) result = await transformStreamingResponse(response)
@@ -146,13 +206,25 @@ async function transformResponseWithThinking(
result = await transformResponse(response) result = await transformResponse(response)
} }
// Apply thinking extraction for high-thinking models
if (!streaming && shouldIncludeThinking(modelName)) {
try { try {
const text = await result.response.clone().text() const text = await result.response.clone().text()
const parsed = JSON.parse(text) as Record<string, unknown>
// Extract and transform thinking blocks if (streaming) {
const signature = extractSignatureFromSsePayload(text)
if (signature) {
setThoughtSignature(fetchInstanceId, signature)
debugLog(`[STREAMING] Stored thought signature for instance ${fetchInstanceId}`)
}
} else {
const parsed = JSON.parse(text) as GeminiResponseBody
const signature = extractSignatureFromResponse(parsed)
if (signature) {
setThoughtSignature(fetchInstanceId, signature)
debugLog(`Stored thought signature for instance ${fetchInstanceId}`)
}
if (shouldIncludeThinking(modelName)) {
const thinkingResult = extractThinkingBlocks(parsed) const thinkingResult = extractThinkingBlocks(parsed)
if (thinkingResult.hasThinking) { if (thinkingResult.hasThinking) {
const transformed = transformResponseThinking(parsed) const transformed = transformResponseThinking(parsed)
@@ -162,10 +234,9 @@ async function transformResponseWithThinking(
headers: result.response.headers, headers: result.response.headers,
}) })
} }
} catch {
// If thinking extraction fails, return original transformed response
} }
} }
} catch {}
return result.response return result.response
} }
@@ -207,9 +278,9 @@ export function createAntigravityFetch(
clientId?: string, clientId?: string,
clientSecret?: string clientSecret?: string
): (url: string, init?: RequestInit) => Promise<Response> { ): (url: string, init?: RequestInit) => Promise<Response> {
// Cache for current token state
let cachedTokens: AntigravityTokens | null = null let cachedTokens: AntigravityTokens | null = null
let cachedProjectId: string | null = null let cachedProjectId: string | null = null
const fetchInstanceId = crypto.randomUUID()
return async (url: string, init: RequestInit = {}): Promise<Response> => { return async (url: string, init: RequestInit = {}): Promise<Response> => {
debugLog(`Intercepting request to: ${url}`) debugLog(`Intercepting request to: ${url}`)
@@ -304,18 +375,23 @@ export function createAntigravityFetch(
} }
const maxEndpoints = Math.min(ANTIGRAVITY_ENDPOINT_FALLBACKS.length, 3) const maxEndpoints = Math.min(ANTIGRAVITY_ENDPOINT_FALLBACKS.length, 3)
const sessionId = getOrCreateSessionId(fetchInstanceId)
const thoughtSignature = getThoughtSignature(fetchInstanceId)
debugLog(`[TSIG][GET] sessionId=${sessionId}, signature=${thoughtSignature ? thoughtSignature.substring(0, 20) + "..." : "none"}`)
for (let i = 0; i < maxEndpoints; i++) { for (let i = 0; i < maxEndpoints; i++) {
const endpoint = ANTIGRAVITY_ENDPOINT_FALLBACKS[i] const endpoint = ANTIGRAVITY_ENDPOINT_FALLBACKS[i]
const response = await attemptFetch( const response = await attemptFetch({
endpoint, endpoint,
url, url,
init, init,
cachedTokens.access_token, accessToken: cachedTokens.access_token,
projectId, projectId,
modelName sessionId,
) modelName,
thoughtSignature,
})
if (response === "pass-through") { if (response === "pass-through") {
debugLog("Non-string body detected, passing through with auth headers") debugLog("Non-string body detected, passing through with auth headers")
@@ -328,7 +404,12 @@ export function createAntigravityFetch(
if (response) { if (response) {
debugLog(`Success with endpoint: ${endpoint}`) debugLog(`Success with endpoint: ${endpoint}`)
return transformResponseWithThinking(response, modelName || "") const transformedResponse = await transformResponseWithThinking(
response,
modelName || "",
fetchInstanceId
)
return transformedResponse
} }
} }

View File

@@ -1,5 +1,3 @@
// Antigravity auth module barrel export
export * from "./types" export * from "./types"
export * from "./constants" export * from "./constants"
export * from "./oauth" export * from "./oauth"
@@ -9,5 +7,6 @@ export * from "./request"
export * from "./response" export * from "./response"
export * from "./tools" export * from "./tools"
export * from "./thinking" export * from "./thinking"
export * from "./thought-signature-store"
export * from "./fetch" export * from "./fetch"
export * from "./plugin" export * from "./plugin"

View File

@@ -132,15 +132,11 @@ function generateRequestId(): string {
return `agent-${crypto.randomUUID()}` return `agent-${crypto.randomUUID()}`
} }
function generateSessionId(): string {
const n = Math.floor(Math.random() * 9_000_000_000_000_000_000)
return `-${n}`
}
export function wrapRequestBody( export function wrapRequestBody(
body: Record<string, unknown>, body: Record<string, unknown>,
projectId: string, projectId: string,
modelName: string modelName: string,
sessionId: string
): AntigravityRequestBody { ): AntigravityRequestBody {
const requestPayload = { ...body } const requestPayload = { ...body }
delete requestPayload.model delete requestPayload.model
@@ -152,11 +148,69 @@ export function wrapRequestBody(
requestId: generateRequestId(), requestId: generateRequestId(),
request: { request: {
...requestPayload, ...requestPayload,
sessionId: generateSessionId(), sessionId,
}, },
} }
} }
interface ContentPart {
functionCall?: Record<string, unknown>
thoughtSignature?: string
[key: string]: unknown
}
interface ContentBlock {
role?: string
parts?: ContentPart[]
[key: string]: unknown
}
function debugLog(message: string): void {
if (process.env.ANTIGRAVITY_DEBUG === "1") {
console.log(`[antigravity-request] ${message}`)
}
}
export function injectThoughtSignatureIntoFunctionCalls(
body: Record<string, unknown>,
signature: string | undefined
): Record<string, unknown> {
debugLog(`[TSIG][INJECT] signature=${signature ? signature.substring(0, 20) + "..." : "none"}`)
if (!signature) {
return body
}
const contents = body.contents as ContentBlock[] | undefined
if (!contents || !Array.isArray(contents)) {
debugLog(`[TSIG][INJECT] no contents array found`)
return body
}
let injectedCount = 0
const modifiedContents = contents.map((content) => {
if (!content.parts || !Array.isArray(content.parts)) {
return content
}
const modifiedParts = content.parts.map((part) => {
if (part.functionCall && !part.thoughtSignature) {
injectedCount++
return {
...part,
thoughtSignature: signature,
}
}
return part
})
return { ...content, parts: modifiedParts }
})
debugLog(`[TSIG][INJECT] injected signature into ${injectedCount} functionCall(s)`)
return { ...body, contents: modifiedContents }
}
/** /**
* Detect if request is for streaming. * Detect if request is for streaming.
* Checks both action name and request body for stream flag. * Checks both action name and request body for stream flag.
@@ -183,48 +237,45 @@ export function isStreamingRequest(
return false return false
} }
/** export interface TransformRequestOptions {
* Transform an OpenAI-format request to Antigravity format. url: string
* This is the main transformation function used by the fetch interceptor. body: Record<string, unknown>
* accessToken: string
* @param url - Original request URL projectId: string
* @param body - Original request body (OpenAI format) sessionId: string
* @param accessToken - OAuth access token for Authorization modelName?: string
* @param projectId - GCP project ID for wrapper
* @param modelName - Model name to use (overrides body.model if provided)
* @param endpointOverride - Optional endpoint override (uses first fallback if not provided)
* @returns Transformed request with URL, headers, body, and streaming flag
*/
export function transformRequest(
url: string,
body: Record<string, unknown>,
accessToken: string,
projectId: string,
modelName?: string,
endpointOverride?: string endpointOverride?: string
): TransformedRequest { thoughtSignature?: string
// Determine model name (parameter override > body > URL) }
export function transformRequest(options: TransformRequestOptions): TransformedRequest {
const {
url,
body,
accessToken,
projectId,
sessionId,
modelName,
endpointOverride,
thoughtSignature,
} = options
const effectiveModel = const effectiveModel =
modelName || extractModelFromBody(body) || extractModelFromUrl(url) || "gemini-3-pro-preview" modelName || extractModelFromBody(body) || extractModelFromUrl(url) || "gemini-3-pro-preview"
// Determine if streaming
const streaming = isStreamingRequest(url, body) const streaming = isStreamingRequest(url, body)
// Determine action (default to appropriate generate action)
const action = streaming ? "streamGenerateContent" : "generateContent" const action = streaming ? "streamGenerateContent" : "generateContent"
// Build URL
const endpoint = endpointOverride || getDefaultEndpoint() const endpoint = endpointOverride || getDefaultEndpoint()
const transformedUrl = buildAntigravityUrl(endpoint, action, streaming) const transformedUrl = buildAntigravityUrl(endpoint, action, streaming)
// Build headers
const headers = buildRequestHeaders(accessToken) const headers = buildRequestHeaders(accessToken)
if (streaming) { if (streaming) {
headers["Accept"] = "text/event-stream" headers["Accept"] = "text/event-stream"
} }
// Wrap body in Antigravity format const bodyWithSignature = injectThoughtSignatureIntoFunctionCalls(body, thoughtSignature)
const wrappedBody = wrapRequestBody(body, projectId, effectiveModel) const wrappedBody = wrapRequestBody(bodyWithSignature, projectId, effectiveModel, sessionId)
return { return {
url: transformedUrl, url: transformedUrl,

View File

@@ -465,6 +465,61 @@ export function isStreamingResponse(response: Response): boolean {
return contentType.includes("text/event-stream") return contentType.includes("text/event-stream")
} }
/**
* Extract thought signature from SSE payload text
*
* Looks for thoughtSignature in SSE events:
* data: { "response": { "candidates": [{ "content": { "parts": [{ "thoughtSignature": "..." }] } }] } }
*
* Returns the last found signature (most recent in the stream).
*
* @param payload - SSE payload text
* @returns Last thought signature if found
*/
export function extractSignatureFromSsePayload(payload: string): string | undefined {
const lines = payload.split("\n")
let lastSignature: string | undefined
for (const line of lines) {
if (!line.startsWith("data:")) {
continue
}
const json = line.slice(5).trim()
if (!json || json === "[DONE]") {
continue
}
try {
const parsed = JSON.parse(json) as Record<string, unknown>
// Check in response wrapper (Antigravity format)
const response = (parsed.response || parsed) as Record<string, unknown>
const candidates = response.candidates as Array<Record<string, unknown>> | undefined
if (candidates && Array.isArray(candidates)) {
for (const candidate of candidates) {
const content = candidate.content as Record<string, unknown> | undefined
const parts = content?.parts as Array<Record<string, unknown>> | undefined
if (parts && Array.isArray(parts)) {
for (const part of parts) {
const sig = (part.thoughtSignature || part.thought_signature) as string | undefined
if (sig && typeof sig === "string") {
lastSignature = sig
}
}
}
}
}
} catch {
// Continue to next line if parsing fails
}
}
return lastSignature
}
/** /**
* Extract usage from SSE payload text * Extract usage from SSE payload text
* *

View File

@@ -0,0 +1,97 @@
/**
* Thought Signature Store
*
* Stores and retrieves thought signatures for multi-turn conversations.
* Gemini 3 Pro requires thought_signature on function call content blocks
* in subsequent requests to maintain reasoning continuity.
*
* Key responsibilities:
* - Store the latest thought signature per session
* - Provide signature for injection into function call requests
* - Clear signatures when sessions end
*/
/**
* In-memory store for thought signatures indexed by session ID
*/
const signatureStore = new Map<string, string>()
/**
* In-memory store for session IDs per fetch instance
* Used to maintain consistent sessionId across multi-turn conversations
*/
const sessionIdStore = new Map<string, string>()
/**
* Store a thought signature for a session
*
* @param sessionKey - Unique session identifier (typically fetch instance ID)
* @param signature - The thought signature from model response
*/
export function setThoughtSignature(sessionKey: string, signature: string): void {
if (sessionKey && signature) {
signatureStore.set(sessionKey, signature)
}
}
/**
* Retrieve the stored thought signature for a session
*
* @param sessionKey - Unique session identifier
* @returns The stored signature or undefined if not found
*/
export function getThoughtSignature(sessionKey: string): string | undefined {
return signatureStore.get(sessionKey)
}
/**
* Clear the thought signature for a session
*
* @param sessionKey - Unique session identifier
*/
export function clearThoughtSignature(sessionKey: string): void {
signatureStore.delete(sessionKey)
}
/**
* Store or retrieve a persistent session ID for a fetch instance
*
* @param fetchInstanceId - Unique identifier for the fetch instance
* @param sessionId - Optional session ID to store (if not provided, returns existing or generates new)
* @returns The session ID for this fetch instance
*/
export function getOrCreateSessionId(fetchInstanceId: string, sessionId?: string): string {
if (sessionId) {
sessionIdStore.set(fetchInstanceId, sessionId)
return sessionId
}
const existing = sessionIdStore.get(fetchInstanceId)
if (existing) {
return existing
}
const n = Math.floor(Math.random() * Number.MAX_SAFE_INTEGER)
const newSessionId = `-${n}`
sessionIdStore.set(fetchInstanceId, newSessionId)
return newSessionId
}
/**
* Clear the session ID for a fetch instance
*
* @param fetchInstanceId - Unique identifier for the fetch instance
*/
export function clearSessionId(fetchInstanceId: string): void {
sessionIdStore.delete(fetchInstanceId)
}
/**
* Clear all stored data for a fetch instance (signature + session ID)
*
* @param fetchInstanceId - Unique identifier for the fetch instance
*/
export function clearFetchInstanceData(fetchInstanceId: string): void {
signatureStore.delete(fetchInstanceId)
sessionIdStore.delete(fetchInstanceId)
}