feat(client): migrate to SessionHandler for per-session project identity#157
Conversation
Replace stateless Handler with muxcore SessionHandler. Each CC session
now gets its own project identity from ProjectContext.Cwd instead of
sharing the daemon's startup project.
Changes:
- engramHandler struct with sync.Map pools (gRPC connections, slug cache)
- HandleRequest: per-session env lookup, project resolution, gRPC dispatch
- OnProjectConnect/OnProjectDisconnect lifecycle logging
- envOrDefault: session env → os.Getenv fallback chain
- resolveProject: cached proxy.ResolveProjectSlug per session
- getOrDialGRPC: connection pool with LoadOrStore pattern
- Simplified main(): no more ResolveProjectSlug(".") or os.Setenv
- Removed: mcpHandler, runTranslator, bufio/io imports
Fixes: cross-project contamination where all daemon sessions shared
one project ID (the daemon's startup cwd).
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
✅ Files skipped from review due to trivial changes (1)
WalkthroughЗаменена потоковая per-session stdin/stdout JSON-RPC обработка на muxcore-native SessionHandler; добавлены хуки жизненного цикла проекта, session‑aware разрешение конфигурации/slug, кэш слагов и пул повторно используемых gRPC‑соединений, запросы теперь обрабатываются как одиночные JSON‑RPC payload'ы. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as JSON‑RPC Client
participant Handler as engramHandler
participant Pool as gRPC Conn Pool
participant GRPC as gRPC Server
Client->>Handler: HandleRequest(ctx, ProjectContext, request []byte)
Handler->>Handler: resolveProject(p) / envOrDefault(p.Env, ...)
Handler->>Pool: Load connKey{addr, tlsMode}
alt conn exists
Pool-->>Handler: existing *grpc.ClientConn
else
Handler->>GRPC: dialGRPC(addr, tlsMode)
GRPC-->>Handler: new *grpc.ClientConn
Handler->>Pool: Store connKey
end
Handler->>GRPC: handleJSONRPC(conn, request)
GRPC-->>Handler: response []byte
Handler-->>Client: response []byte, error
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the engram entry point to implement the muxcore.SessionHandler and muxcore.ProjectLifecycle interfaces, introducing gRPC connection pooling and project slug caching. Feedback highlights a critical need to remove a local file system replace directive in go.mod that would break builds in other environments. Furthermore, the gRPC connection pooling logic requires the API token to be part of the cache key to ensure proper authentication isolation, and the project slug cache should be pruned during disconnection to prevent memory leaks.
|
|
||
| replace github.com/thebtf/mcp-mux/muxcore v0.18.2 => D:\Dev\mcp-mux\muxcore |
There was a problem hiding this comment.
The local replace directive pointing to an absolute Windows path (D:\Dev\mcp-mux\muxcore) will break builds for other developers and in CI environments. This should be removed before merging.
| type connKey struct { | ||
| addr string | ||
| tlsMode string // "custom-ca", "system-tls", "plaintext" | ||
| } |
There was a problem hiding this comment.
The connKey should include the token because the gRPC connection is authenticated via a tokenInterceptor that is fixed at dial time. Without the token in the key, sessions with different API tokens but the same server URL will incorrectly share the same authenticated connection, leading to potential security issues or incorrect project access.
| type connKey struct { | |
| addr string | |
| tlsMode string // "custom-ca", "system-tls", "plaintext" | |
| } | |
| type connKey struct { | |
| addr string | |
| tlsMode string // "custom-ca", "system-tls", "plaintext" | |
| token string | |
| } |
| tlsMode = "system-tls" | ||
| } | ||
|
|
||
| key := connKey{addr: grpcAddr, tlsMode: tlsMode} |
There was a problem hiding this comment.
Include the token in the connKey to ensure that connections are pooled correctly based on their authentication credentials.
| key := connKey{addr: grpcAddr, tlsMode: tlsMode} | |
| key := connKey{addr: grpcAddr, tlsMode: tlsMode, token: token} |
| func (h *engramHandler) OnProjectDisconnect(projectID string) { | ||
| fmt.Fprintf(os.Stderr, "[engram] session disconnected: project=%s\n", projectID) | ||
| } |
There was a problem hiding this comment.
The slugCache should be cleaned up when a project disconnects to prevent an unbounded memory leak as new sessions are created over time. Additionally, logging the resolved slug (if available) provides better diagnostic information than the internal project ID.
| func (h *engramHandler) OnProjectDisconnect(projectID string) { | |
| fmt.Fprintf(os.Stderr, "[engram] session disconnected: project=%s\n", projectID) | |
| } | |
| func (h *engramHandler) OnProjectDisconnect(projectID string) { | |
| slug := projectID | |
| if cached, ok := h.slugCache.LoadAndDelete(projectID); ok { | |
| slug = cached.(string) | |
| } | |
| fmt.Fprintf(os.Stderr, "[engram] session disconnected: project=%s\n", slug) | |
| } |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
cmd/engram/main_test.go (2)
47-53: Используйте уникальный ключ для изоляции теста.
os.Unsetenvне восстанавливает состояние окружения после теста. Хотя ключ уникальный (MISSING_KEY_TEST_12345), это хорошая практика — убедиться, что тест изолирован.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/engram/main_test.go` around lines 47 - 53, Test TestEnvOrDefault_NeitherSet uses os.Unsetenv but doesn't restore the original environment value; update the test to save the prior value of "MISSING_KEY_TEST_12345" using os.LookupEnv, then unset for the test and defer restoring the environment (either os.Setenv with the saved value or os.Unsetenv if it was absent) so envOrDefault is tested in isolation.
85-104: Тест не проверяет значимую функциональность.Строка
_ = hничего не проверяет. Комментарий утверждает, что проверяется "usability" handler'а, но фактически тест только сравнивает struct'ыconnKey. Рассмотрите возможность удаления неиспользуемого handler'а или добавления реальной проверки.♻️ Предлагаемый рефакторинг
func TestGRPCPool_SharedConnection(t *testing.T) { - h := &engramHandler{} - // Two calls with the same URL should return the same connection. // We can't easily test real gRPC connections without a server, // but we can verify the pool key logic. key1 := connKey{addr: "host:37777", tlsMode: "plaintext"} key2 := connKey{addr: "host:37777", tlsMode: "plaintext"} if key1 != key2 { t.Error("identical connKeys should be equal") } key3 := connKey{addr: "other:443", tlsMode: "system-tls"} if key1 == key3 { t.Error("different connKeys should not be equal") } - - // Verify the handler struct is usable (sync.Map initialized by zero value). - _ = h }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/engram/main_test.go` around lines 85 - 104, The test currently creates engramHandler h and then does `_ = h` which asserts nothing; either remove the unused engramHandler variable or add a real assertion that its zero-valued sync.Map is usable: replace `_ = h` with a small exercise of engramHandler (e.g., call an existing method that uses its internal sync.Map or, if no method exists, add a temporary store/load via the handler's connection map field) to store a value and read it back to confirm zero-value initialization; reference the engramHandler and TestGRPCPool_SharedConnection identifiers when making the change.cmd/engram/main.go (1)
103-107: Несогласованность логирования:projectIDvs cached slug.
OnProjectConnectлогирует resolved slug, аOnProjectDisconnectлогирует rawprojectID. Для согласованности можно использовать cached slug изslugCache:♻️ Предлагаемый рефакторинг
// OnProjectDisconnect is called when a CC session disconnects. // Implements muxcore.ProjectLifecycle. func (h *engramHandler) OnProjectDisconnect(projectID string) { - fmt.Fprintf(os.Stderr, "[engram] session disconnected: project=%s\n", projectID) + slug := projectID + if cached, ok := h.slugCache.Load(projectID); ok { + slug = cached.(string) + } + fmt.Fprintf(os.Stderr, "[engram] session disconnected: project=%s\n", slug) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/engram/main.go` around lines 103 - 107, OnProjectDisconnect currently logs the raw projectID while OnProjectConnect logs the resolved slug; update engramHandler.OnProjectDisconnect to look up the cached slug from h.slugCache (the same cache used by OnProjectConnect) and log that slug for consistency, falling back to the original projectID if no slug is found; modify the logging in OnProjectDisconnect to use the cached value and include clear context (e.g., "session disconnected: project=<slug|projectID>").
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cmd/engram/main.go`:
- Around line 82-90: When json.Unmarshal into jsonrpcRequest fails, req.ID may
be invalid or partial; update the error branch to explicitly return a
jsonrpcResponse with JSONRPC "2.0", ID set to nil (null in output) and the parse
error (jsonrpcError Code -32700) instead of using req.ID. Locate the
unmarshalling block around json.Unmarshal(request, &req) and change the created
jsonrpcResponse to use an explicit nil ID value rather than req.ID so the
response follows JSON-RPC 2.0 (null id) on parse errors.
In `@go.mod`:
- Around line 64-65: Удалите локальную директиву replace в go.mod, которая
указывает на Windows-путь "D:\Dev\mcp-mux\muxcore" (строка с replace
github.com/thebtf/mcp-mux/muxcore v0.18.2 => D:\Dev\mcp-mux\muxcore), чтобы
сборка CI и другие разработчики не ломались; либо замените её на корректную
версия-модуль (например вернуть на v0.18.2) или удалите строку полностью перед
мержем.
---
Nitpick comments:
In `@cmd/engram/main_test.go`:
- Around line 47-53: Test TestEnvOrDefault_NeitherSet uses os.Unsetenv but
doesn't restore the original environment value; update the test to save the
prior value of "MISSING_KEY_TEST_12345" using os.LookupEnv, then unset for the
test and defer restoring the environment (either os.Setenv with the saved value
or os.Unsetenv if it was absent) so envOrDefault is tested in isolation.
- Around line 85-104: The test currently creates engramHandler h and then does
`_ = h` which asserts nothing; either remove the unused engramHandler variable
or add a real assertion that its zero-valued sync.Map is usable: replace `_ = h`
with a small exercise of engramHandler (e.g., call an existing method that uses
its internal sync.Map or, if no method exists, add a temporary store/load via
the handler's connection map field) to store a value and read it back to confirm
zero-value initialization; reference the engramHandler and
TestGRPCPool_SharedConnection identifiers when making the change.
In `@cmd/engram/main.go`:
- Around line 103-107: OnProjectDisconnect currently logs the raw projectID
while OnProjectConnect logs the resolved slug; update
engramHandler.OnProjectDisconnect to look up the cached slug from h.slugCache
(the same cache used by OnProjectConnect) and log that slug for consistency,
falling back to the original projectID if no slug is found; modify the logging
in OnProjectDisconnect to use the cached value and include clear context (e.g.,
"session disconnected: project=<slug|projectID>").
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 83dd43aa-0b1b-4d1c-876d-935f38ea9011
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (3)
cmd/engram/main.gocmd/engram/main_test.gogo.mod
| var req jsonrpcRequest | ||
| if err := json.Unmarshal(request, &req); err != nil { | ||
| resp := jsonrpcResponse{ | ||
| JSONRPC: "2.0", | ||
| ID: req.ID, | ||
| Error: &jsonrpcError{Code: -32700, Message: "Parse error"}, | ||
| } | ||
| return json.Marshal(resp) | ||
| } |
There was a problem hiding this comment.
При ошибке парсинга req.ID может быть невалидным.
Если json.Unmarshal завершается с ошибкой, req.ID может остаться nil или содержать частичные данные. По спецификации JSON-RPC 2.0, при ошибке парсинга id должен быть null.
🐛 Предлагаемое исправление
var req jsonrpcRequest
if err := json.Unmarshal(request, &req); err != nil {
resp := jsonrpcResponse{
JSONRPC: "2.0",
- ID: req.ID,
+ ID: nil, // JSON-RPC spec: id MUST be null for parse errors
Error: &jsonrpcError{Code: -32700, Message: "Parse error"},
}
return json.Marshal(resp)
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| var req jsonrpcRequest | |
| if err := json.Unmarshal(request, &req); err != nil { | |
| resp := jsonrpcResponse{ | |
| JSONRPC: "2.0", | |
| ID: req.ID, | |
| Error: &jsonrpcError{Code: -32700, Message: "Parse error"}, | |
| } | |
| return json.Marshal(resp) | |
| } | |
| var req jsonrpcRequest | |
| if err := json.Unmarshal(request, &req); err != nil { | |
| resp := jsonrpcResponse{ | |
| JSONRPC: "2.0", | |
| ID: nil, // JSON-RPC spec: id MUST be null for parse errors | |
| Error: &jsonrpcError{Code: -32700, Message: "Parse error"}, | |
| } | |
| return json.Marshal(resp) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cmd/engram/main.go` around lines 82 - 90, When json.Unmarshal into
jsonrpcRequest fails, req.ID may be invalid or partial; update the error branch
to explicitly return a jsonrpcResponse with JSONRPC "2.0", ID set to nil (null
in output) and the parse error (jsonrpcError Code -32700) instead of using
req.ID. Locate the unmarshalling block around json.Unmarshal(request, &req) and
change the created jsonrpcResponse to use an explicit nil ID value rather than
req.ID so the response follows JSON-RPC 2.0 (null id) on parse errors.
1 participant
Summary
Migrates the engram client from stateless
Handlerto muxcoreSessionHandlerfor per-session project identity. Fixes the critical bug where all CC sessions through the daemon shared one project ID.Architecture doc:
.agent/specs/session-architecture/architecture.mdSpec:
.agent/specs/session-architecture/spec.md(7 FR, 4 NFR, 4 US)Depends on: muxcore/v0.18.2 (PR thebtf/mcp-mux#49, merged)
Changes
cmd/engram/main.go
muxcore.SessionHandler+ProjectLifecycleproxy.ResolveProjectSlug(p.Cwd)per sessionsync.MapLoadOrStore patternResolveProjectSlug(".")oros.SetenvmcpHandler,runTranslator, unused importscmd/engram/main_test.go (new)
go.mod
What this fixes
Before: all sessions through daemon get project=engram (daemon startup cwd)
After: Session A (cwd: nvmdfs) gets project=nvmdfs, Session B (cwd: engram) gets project=engram
Summary by CodeRabbit
Примечания к выпуску
Улучшения производительности
Улучшения надёжности
Новые возможности
Тесты
Обновления