Token storage chunking fixes - #22
Conversation
…token_storage_chunking
WalkthroughRefactors CLI token storage to support dynamic multi-chunk tokens with a helper for chunk count retrieval, updates GetRawToken to concatenate chunks with a legacy fallback, adjusts SetRawToken for dynamic cleanup on save/delete, removes obsolete parsing, and updates one test’s expected error message. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant CliSession
participant Keyring
Caller->>CliSession: GetRawToken()
CliSession->>Keyring: Read chunk_count
alt chunk_count available and parsed
loop for i in 0..count-1
CliSession->>Keyring: Read chunk_i
end
CliSession->>CliSession: Concatenate chunks
CliSession->>CliSession: Unmarshal oauth2.Token
CliSession-->>Caller: Token
else fallback
CliSession->>Keyring: Read legacy token key
CliSession->>CliSession: Unmarshal oauth2.Token
CliSession-->>Caller: Token or error
end
sequenceDiagram
participant Caller
participant CliSession
participant Keyring
Caller->>CliSession: SetRawToken(token or nil)
alt token == nil (delete)
CliSession->>Keyring: Read chunk_count (if exists)
loop for i in 0..count-1
CliSession->>Keyring: Delete chunk_i
end
CliSession->>Keyring: Delete chunk_count
CliSession-->>Caller: Done
else save token
CliSession->>Keyring: Read chunk_count (if exists)
loop for i in 0..count-1
CliSession->>Keyring: Delete chunk_i
end
CliSession->>CliSession: Marshal token bytes and chunk
loop for each new chunk_i
CliSession->>Keyring: Write chunk_i
end
CliSession->>Keyring: Write chunk_count
CliSession-->>Caller: Done
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15 minutes Possibly related PRs
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🔭 Outside diff range comments (1)
frameworks/cli/cliSession.go (1)
112-124: Fix compile errors in chunk-saving loop (range over int, undefined min).
for i := range chunksdoes not compile becausechunksis anint.minis not defined anywhere in this package.Apply this diff to fix both issues:
- for i := range chunks { - start := i * chunkSize - end := min(start+chunkSize, len(t)) + for i := 0; i < chunks; i++ { + start := i * chunkSize + end := start + chunkSize + if end > len(t) { + end = len(t) + }
🧹 Nitpick comments (6)
frameworks/cli/cliSession.go (5)
101-110: Avoid shadowing and improve cleanup robustness when removing existing chunks.Shadowing
chunksmakes the code harder to read. Also, breaking on the first removal error may leave stale keys behind.Apply this diff:
- // Try to get chunk count - if chunks, err := c.getChunkCount(countKey); err == nil { - // Remove any existing chunks - for i := 0; i < chunks; i++ { + // Try to get existing chunk count for cleanup + if prevChunks, err := c.getChunkCount(countKey); err == nil { + // Remove any existing chunks (best-effort) + for i := 0; i < prevChunks; i++ { chunkKey := fmt.Sprintf("%s_chunk_%d", key, i) - if err := c.keyring.Remove(chunkKey); err != nil { - break - } + _ = c.keyring.Remove(chunkKey) // ignore errors to maximize cleanup } }
81-87: Use best-effort cleanup during deletion instead of breaking on first error.Breaking on the first error can leave leftover chunks; prefer continuing the loop.
- if err := c.keyring.Remove(chunkKey); err != nil { - break - } + _ = c.keyring.Remove(chunkKey) // ignore errors to continue best-effort cleanup
22-33: Parsing chunk count works, but consider using strconv and validating > 0.
fmt.Sscanfis fine, butstrconv.Atoi(strings.TrimSpace(...))is simpler and avoids partial matches. Also consider rejecting non-positive counts to prevent empty concatenations.Example diff (plus imports for strconv/strings):
- var chunks int - if _, err := fmt.Sscanf(string(countItem.Data), "%d", &chunks); err != nil { - return 0, fmt.Errorf("failed to parse chunk count: %w", err) - } + s := strings.TrimSpace(string(countItem.Data)) + chunks, err := strconv.Atoi(s) + if err != nil { + return 0, fmt.Errorf("failed to parse chunk count: %w", err) + } + if chunks <= 0 { + return 0, fmt.Errorf("invalid chunk count: %d", chunks) + }Outside this hunk, add imports:
import ( "encoding/json" "fmt" "strconv" "strings" // ... )
126-126: Redundant re-assignment of countKey.
countKeyis already defined earlier; reassigning it to the same computed value adds noise.- countKey = fmt.Sprintf("%s_chunk_count", key)
73-136: Optional: also write the legacy single-key token to preserve fallback behavior.Currently, saving writes only chunked data. If the chunk count key ever becomes unreadable, the fallback will look for the legacy
"kinde_token", which could be stale or missing. Consider also setting the legacy key with the full JSON for resilience:// After saving chunks and chunk count: _ = c.keyring.Set(keyring.Item{Key: key, Data: t}) // ignore error or handle as neededOr, if you prefer to avoid duplicating storage, explicitly remove the legacy key during save to prevent stale fallback:
_ = c.keyring.Remove(key)frameworks/cli/cliSession_test.go (1)
113-113: Relax brittle assertion to avoid coupling to keyring’s exact error string.Asserting the full propagated message ties the test to keyring’s wording. Prefer asserting the high-level prefix and/or using
assert.ErrorIs.- assert.Contains(err.Error(), "failed to get token: The specified item could not be found in the keyring") + assert.Contains(err.Error(), "failed to get token") + assert.ErrorIs(err, keyring.ErrKeyNotFound)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (2)
frameworks/cli/cliSession.go(4 hunks)frameworks/cli/cliSession_test.go(1 hunks)
🔇 Additional comments (2)
frameworks/cli/cliSession.go (2)
41-53: LGTM: Fallback to legacy token on chunk-count error is sensible.Gracefully falling back to the legacy single-key token when chunked storage isn't available or count parsing fails is a good compatibility choice and matches the updated test.
127-133: No action needed — repo declares Go >= 1.20 (fmt.Appendf is supported)Verified that the repository's go directives are >= 1.20 (root go.mod and examples declare go 1.24.4), so fmt.Appendf usage is supported.
- go.mod files checked: go.mod, examples/cli/go.mod, examples/gin-chat/go.mod — all: go 1.24.4
- Files using fmt.Appendf:
- frameworks/cli/cliSession.go (countData := fmt.Appendf(nil, "%d", chunks) at ~line 127)
- kinde/kinde_test.go (test usage at line 78)
|
closing as some commits ended-up not signed |
1 participant
Explain your changes
Fixed issues with unbound counts for token chinks