Skip to content

token cache improvements - #9652

Closed
pjain1 wants to merge 2 commits into
mainfrom
token_caching_improvements
Closed

pjain1 wants to merge 2 commits into
mainfrom
token_caching_improvements

Conversation

@pjain1

@pjain1 pjain1 commented Jul 6, 2026

Copy link
Copy Markdown
Member
  1. Token based truncation of messages in steps using max_input_tokens configurable per connector.
  2. Split the analyst prompt into a session-stable systemPrompt and a per-turn contextPrompt (date, dashboard state, filters)
  3. prompt_cache_key for per session cache for openai connector.

Checklist:

  • Covered by tests
  • Ran it and it works as intended
  • Reviewed the diff before requesting a review
  • Checked for unhandled edge cases
  • Linked the issues it closes
  • Checked if the docs need to be updated. If so, create a separate Linear DOCS issue
  • Intend to cherry-pick into the release branch
  • I'm proud of this work!

@nishantmonu51 nishantmonu51 added Type:Improvement Size:M Medium change: 100-499 lines labels Jul 6, 2026
@nishantmonu51

Copy link
Copy Markdown
Collaborator

@codex: review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c3be0511c0

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

// Insert the context prompt just before the trailing user prompt. It is regenerated on every
// call and not stored in the session, so keeping it at the end preserves the cached prefix of
// the earlier messages.
messages = slices.Insert(messages, len(messages)-1, NewTextCompletionMessage(RoleUser, contextPrompt))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Keep admin AI instructions in a system message

When a project defines ai_instructions, this inserts the new contextPrompt as a user message immediately before the actual user prompt, and that prompt now contains the administrator project-wide instructions. That downgrades those instructions from the previous system prompt to user-level content that the following user message can contradict or override, so configured analyst behavior/guardrails can be ignored. Keep admin instructions in the system prompt or another higher-priority message while moving only truly dynamic dashboard context here.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

handled

@nishantmonu51

Copy link
Copy Markdown
Collaborator

@codex: review again

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bfad59c65f

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

// Insert the context prompt just before the trailing user prompt. It is regenerated on every
// call and not stored in the session, so keeping it at the end preserves the cached prefix of
// the earlier messages.
messages = slices.Insert(messages, len(messages)-1, NewTextCompletionMessage(RoleUser, contextPrompt))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve per-turn context during truncation

When an analyst completion runs long enough for token truncation after tool calls have been appended, this inserted context message is no longer near the tail; maybeTruncateMessages only preserves the first four messages and a suffix, so it can drop the current dashboard filters/time range/date while keeping later tool-call history. In large dashboard sessions, later iterations can therefore query or answer without the per-turn context that previously lived in the always-kept system prompt.

Useful? React with 👍 / 👎.

@nishantmonu51 nishantmonu51 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The branch is 209 commits behind main and git merge-tree reports content conflicts in both runtime/ai/ai.go and runtime/ai/analyst_agent.go, so a rebase is needed before this is mergeable.

Two smaller things:

  • The truncation indicator is emitted with role system, and the Claude driver's convertMessages lifts every system message into the system block, so on Claude the "[N messages omitted]" text mutates the cached system prefix each time the skip count steps. This is pre-existing on main, but it works directly against this PR's caching goal.
  • The admin driver's new max_input_tokens spec property is unreachable: its config is built from a fixed map in admin/deployments.go:270 and cli/pkg/local/app.go:281 (admin_url, access_token, project_id), so only the 200k default can ever apply.

Comment thread runtime/ai/ai.go
Comment on lines +1648 to 1656
for keptTokens > maxTokens && truncateKeepFirst+skipped < len(messages)-1 {
keptTokens -= est[truncateKeepFirst+skipped]
skipped++
}

if len(messages) <= maxMessages {
if skipped == 0 {
return messages
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Once skipped goes above zero, the round-up to truncateStep and the len(messages)-truncateKeepFirst-1 cap together keep only the final message for any conversation shorter than about 25 messages, so a single turn can lose the user's question, the context prompt and every tool result. On the first turn of an explore the seeded messages are [system, call, result, call, result, context, user] and the loop then appends (call, result) pairs of up to maxMessageSizeBytes each; query results at the default 250-row limit run 30-60KB, so proto.Size/3 crosses the 136k default budget after roughly nine iterations, at len = 25. The minimal skip of 1 rounds up to 20, the cap keeps only index 24, and the unbalanced-pair cleanup then drops the seed call at index 3 and that last result, leaving the model [system, seed call, seed result, "[20 messages omitted]"]. main cannot reach this state because count-based truncation at maxMessages = 100 with keepLast = 16 never fires within one turn; the step rounding needs a floor that preserves a recent window, and the context and user prompts should be pinned rather than relying on their position.

Comment on lines +177 to +185
// Build completion messages.
// The system prompt contains only instructions that stay stable for the duration of a session;
// per-turn dynamic info (date, dashboard state, etc.) goes in a separate context message near the
// end, so the LLM prompt cache prefix stays valid across turns.
systemPrompt, err := t.systemPrompt(ctx, args)
if err != nil {
return nil, err
}
contextPrompt, err := t.contextPrompt(ctx, metricsViewNames, args)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

main has since landed #9645, which already splits the analyst prompt into a stable systemPrompt() and a per-turn userPrompt(ctx, metricsViewNames, args), with #9851 layered on top, so contextPrompt and the template rewrite here duplicate work that is already merged and should be dropped on rebase rather than reconciled. The slices.Insert(messages, len(messages)-1, ...) on line 207 also does not survive a mechanical rebase: on current main the layout ends [..., user prompt, seeded tool calls], so the last message is a tool result and the insert would land between a seeded tool call and its result.

@pjain1

pjain1 commented Sep 21, 2026

Copy link
Copy Markdown
Member Author

@nishantmonu51 Will close this PR since #9645 handled most of the improvements.

@pjain1 pjain1 closed this Sep 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Size:M Medium change: 100-499 lines Type:Improvement

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants