feat(probe): surface canonical token usage, lift tool calls, prune dead code - #1496
Open
FFengIll wants to merge 9 commits into
Open
feat(probe): surface canonical token usage, lift tool calls, prune dead code#1496FFengIll wants to merge 9 commits into
FFengIll wants to merge 9 commits into
Conversation
Token usage was declared on Result but never populated — every SDK probe JSON-marshaled the whole response into Content and returned, leaving the usage fields at zero. Tool calls were never lifted either. Resolve by reusing the existing internal/protocol/usage extractors rather than re-implementing parsing: - Result now carries a single Usage *protocol.TokenUsage in the canonical shape (input_tokens/output_tokens/cache_read_tokens/cache_write_tokens/ reasoning_tokens) — the same vocabulary the rest of TB emits and the frontend renders. No parallel flat token fields. - Non-stream probes call FromOpenAIChatCompletion/FromOpenAIResponses/ FromAnthropicMessage on resp.Usage. Stream probes take usage from the final chunk (OpenAI, with stream_options.include_usage now requested) / response.completed event (Responses) / AnthropicAccumulator (Anthropic). - Tool mode lifts calls into ToolCalls via per-provider extractors. - Success is now set in toProbeResult (was a never-true bug); LatencyMs is owned by the SDK probe (pure upstream time) — the handler no longer overwrites it. Also remove meaningless/dead implementation surfaced during review: - defaultModelForAPIStyle + model guessing: a probe tests a specific model the user picked; an empty model now errors explicitly instead of silently inventing claude-3-haiku/gpt-3.5-turbo. - ProbeStream + probeProviderStream: 1-line passthroughs; Probe handles all modes (stream decision lives in the SDK helpers via testMode). - ProbeProviderWithSDK -> unexported probeProviderWithSDK (the testProviderConnectivity reuse it was exported for no longer exists); routing-capture wiring deduped via a maybeCapture closure. - Dead types Request, ProviderRequest, ProbeProviderResponseData, and the orphaned ProbeProviderResponse envelope (zero refs/swagger/frontend). Google is out of scope: no usage/tool-call extraction, probes unchanged.
Drop the non-canonical prompt_tokens/completion_tokens/total_tokens fields from the probe result types and read the canonical usage object instead (protocol.TokenUsage shape: input_tokens/output_tokens/cache_read_tokens/ cache_write_tokens). The token chip now derives its total from usage.input_tokens + usage.output_tokens, matching TokenUsage.TotalTokens() on the backend. Updates the MSW mock and the in-dialog demo data to match. Backend model change requires `task codegen` to regenerate schema.d.ts / swagger; the hand-written types/probe.ts is updated as a placeholder.
Refresh the stale code-layout block (result.go/e2e.go/sdkprobe.go/ lightweight.go were all renamed) and document the Result field semantics: canonical protocol.TokenUsage (no parallel flat fields), SDK-owned LatencyMs, lifted tool calls, and the routing-trace fields. Note that probeProviderWithSDK reuses internal/protocol/usage rather than re-implementing parsing.
…enominator The light probe had four near-identical probeXxxEndpoint methods returning parallel endpointReport/modelsReport structs that Probe() then field-copied into LightweightProbeResponseData one by one. Collapse to focused helpers that write outcomes directly into the target fields: - runOptionsEndpoint (HTTP-level OPTIONS, no SDK) - runChatEndpoint/runResponsesEndpoint over a shared runOpenAIEndpoint body (only the SDK call + success label differ) - runModelsEndpoint with a clean switch over error cases Also fixes a latent bug: the summary message hardcoded "/4 endpoints", so non-OpenAI providers (which skip chat/responses) reported e.g. "2/4". The denominator now reflects the endpoints actually run.
…/vmodel paths
The three loopback resolution paths (resolveProviderTarget's loopback half,
resolveRuleTarget, resolveVModelLoopbackTarget) each built the same
provider_config target by hand: {Name, APIBase: loopback, APIStyle,
Token: model token, Model}. Extract loopbackConfigTarget so the loopback
construction lives in one place; each call site now passes only the
name/apiStyle/model that actually differ.
The OpenAI probe tool's function parameters used the JSON key "type:" (colon inside the key) instead of "type", so the schema field was misnamed. Aligns with the other two tool definitions in the same file.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
After the probe rename, the implementation still had the original gaps it inherited: token usage was declared on
Resultbut never actually populated (every probe JSON-marshaled the whole response intoContentand returned zero usage), tool calls were never surfaced, severalResultfields were dead or irregular, and the E2E/Light paths duplicated logic.This PR fixes the substance so the probe returns trustworthy, canonical data the rest of TB and the frontend already understand.
Key Changes
Token usage is real and canonical:
Result.Usageis now a single*protocol.TokenUsagein the canonical shape (input_tokens/output_tokens/cache_read_tokens/cache_write_tokens/reasoning_tokens), populated by reusing the existinginternal/protocol/usageextractors rather than re-implementing parsing. The probe no longer emits a parallelprompt_tokens/completion_tokens/total_tokensvocabulary that diverged from everything else — consumers read one consistent shape. Streaming probesnow surface real usage (OpenAI Chat requests
stream_options.include_usage; Responses readsresponse.completed; Anthropic uses the accumulator). Google is intentionally out of scope.Tool-mode probes show what the model called: tool calls are lifted out of the response into
Result.ToolCalls({id, name, input}) via per-provider extractors, so a tool probe reports the actual function call instead of burying it in the raw response blob.Result fields mean what they say:
Successis now set on success (it was a never-true bug — only the cache-hit path set it);LatencyMsis the pure upstream round-trip time owned by the SDK probe rather than a handler-overwritten end-to-end number. Removeddead surface:
ModelsCounton the E2E result, the unusedE2EResponseChunk, and three orphaned request/response types plus theProbeProviderResponseenvelope that nothing referenced.Connectivity matrix is correct for every provider: the light probe's summary no longer hardcodes "/4 endpoints" — non-OpenAI providers (which skip chat/responses) now report a true denominator instead of "2/4". The four near-identical endpoint helpers were
collapsed into focused ones sharing one body.
Frontend renders the canonical usage: the probe dialog's token chip derives its total from
usage.input_tokens + usage.output_tokens(mirroringTokenUsage.TotalTokens()), and the hand-writtentypes/probe.ts+ mock + demo data follow the canonical shape.Notes
task codegento regenerateschema.d.ts/swagger;types/probe.tsis updated as a placeholder until then.go build ./internal/...,go vet,go test ./internal/probe/...(62 cases) pass; frontend typecheck introduces no new errors (7 pre-existing legacy quota errors unchanged)."type:"typo (colon inside the JSON key) in the OpenAI probe tool's schema that misnamed the field.