Skip to content

feat: add schema-driven reasoning and improve agent reliability - #6

Open
trankien84 wants to merge 5 commits into
zsecducna:mainfrom
trankien84:local-main
Open

feat: add schema-driven reasoning and improve agent reliability#6
trankien84 wants to merge 5 commits into
zsecducna:mainfrom
trankien84:local-main

Conversation

@trankien84

@trankien84 trankien84 commented Jul 22, 2026

Copy link
Copy Markdown

Summary

This PR adds runtime schema-driven native reasoning support and improves reliability for long-running agent workflows.

The main changes are:

  1. Native thinking and reasoning effort support based on each model's runtime schema.
  2. Opt-in diagnostics for reasoning, payloads, streams, and chunk normalization.
  3. Reliability improvements for agent tool execution, long streams, and profile region discovery.

1. Schema-driven native reasoning

Problem

Kiro exposes model-specific request capabilities through additionalModelRequestFieldsSchema, but the proxy did not preserve or use that schema when building requests.

As a result:

  • native thinking and effort fields could not be forwarded reliably;
  • supported effort values and field paths had to be assumed;
  • Claude Messages, OpenAI Chat Completions, and Responses requests behaved differently;
  • the legacy -thinking suffix relied mainly on prompt injection instead of model-native request fields.

Solution

The proxy now preserves additionalModelRequestFieldsSchema from ListAvailableModels and parses it at runtime.

The parsed capability includes:

  • supported thinking types;
  • supported thinking display values;
  • support for budget_tokens;
  • the native effort path:
    • output_config.effort, or
    • reasoning.effort;
  • supported effort values such as low, medium, high, xhigh, and max.

Requested values are validated against the selected model's schema before the upstream request is sent.

API coverage

Native reasoning is supported across:

  • Anthropic Messages API;
  • OpenAI Chat Completions API;
  • OpenAI Responses API;
  • Claude token-count requests.

The generated fields are forwarded through
additionalModelRequestFields.

Backward compatibility

The configured thinking suffix, such as -thinking, remains supported.

When compatible native reasoning fields are available, they are forwarded directly and the legacy thinking prompt is not injected.

When native fields are unavailable, the existing legacy thinking prompt remains the fallback.


2. Environment-gated diagnostics

Problem

Investigating issues such as truncated payloads, incomplete streams, reasoning forwarding, and chunk normalization previously required adding temporary logs directly to the request path.

Solution

This PR adds opt-in diagnostic flags:

Variable Description Default
KIRO_DIAGNOSTICS Enables stream, payload, and reasoning diagnostics false
KIRO_DIAG_STREAM Logs aggregate stream metrics false
KIRO_DIAG_PAYLOAD Logs payload size and truncation metrics false
KIRO_DIAG_REASONING Logs reasoning capabilities and forwarded fields false
KIRO_DIAG_CHUNKS Logs raw and normalized stream chunks; may contain sensitive content false

Accepted enabled values are 1, true, yes, on, and enabled.

KIRO_DIAGNOSTICS intentionally does not enable
KIRO_DIAG_CHUNKS, because chunk logs may contain prompts, reasoning, source code, or tool output.

All diagnostics remain disabled by default.


3. Agent tool-execution reliability

Problem

During coding-agent workflows, a model may describe or plan a tool action without actually invoking the required tool, causing the turn to end before the requested work is complete.

Solution

Tool-enabled agent requests now include execution guidance that instructs the model to:

  • invoke the appropriate tool instead of only describing the action;
  • continue tool execution while required work remains;
  • validate the result before returning the final response.

A shorter reminder is also added on continued tool turns when pending work still requires tool execution.

The guidance is not added when tool use is explicitly disabled through tool_choice.


4. Stream and region reliability

Longer streaming timeout

The Kiro streaming HTTP timeout is increased from 5 minutes to 15 minutes to support long-running coding-agent operations.

The separate REST client keeps its existing 30-second timeout.

IDC and external IdP region probing

For idc and external_idp accounts, the authentication region may differ from the Kiro profile region.

Profile discovery can now probe configured fallback regions for these account types when resolving the profile ARN.

Once a profile ARN is resolved, its region continues to determine the data-plane endpoint.


5. Payload and stream diagnostics

Payload diagnostics report:

  • serialized byte size;
  • estimated token usage;
  • history size;
  • tool and tool-result counts;
  • current-message size;
  • before/after truncation metrics.

The diagnostics wrap the existing byte-and-token-aware truncation logic without changing its preservation rules.

Stream diagnostics report aggregate information such as:

  • assistant events and characters;
  • reasoning events and characters;
  • tool events;
  • token usage;
  • credits;
  • completion reason and duration.

Raw chunk previews remain separately opt-in.


Changes

  • Added a diagnostics package with startup-time environment flags.
  • Preserved model-specific request schemas from ListAvailableModels.
  • Added runtime reasoning-capability parsing and validation.
  • Added native request-field generation for Claude and OpenAI-compatible request shapes.
  • Added Responses API reasoning compatibility.
  • Added outbound reasoning-field forwarding.
  • Retained legacy thinking fallback behavior.
  • Added agent tool-execution guidance.
  • Increased the streaming timeout while preserving the short REST timeout.
  • Added IDC and external IdP profile-region fallback probing.
  • Documented the new diagnostic environment variables in README.md.

Tests

Added and updated tests covering:

  • parsing output_config.effort;
  • parsing reasoning.effort;
  • thinking type discovery;
  • thinking display discovery;
  • supported effort normalization;
  • native Claude request-field generation;
  • rejection of unsupported thinking types;
  • rejection of unsupported display values;
  • rejection of unsupported effort values;
  • separation of streaming and REST timeouts.

Verification performed:

  • gofmt
  • go test ./...
  • go vet ./...
  • git diff --check zsec/main --

Compatibility notes

  • Existing requests without native reasoning fields keep their previous behavior.
  • Existing -thinking model suffixes remain supported.
  • Diagnostics are disabled unless explicitly enabled.
  • Raw chunk logging requires a separate opt-in.
  • The 30-second REST timeout is unchanged.

root and others added 5 commits July 15, 2026 11:38
Normalize Claude thinking.type=enabled to a model-supported adaptive type before schema validation.

Keep strict validation for unsupported thinking types, effort levels, display values, and budget fields, and add regression coverage for Claude and OpenAI-compatible reasoning requests.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant