Skip to content

Commit 06671eb

Browse files
billyb2billyb2
authored
Merge pull request #380 from depot/billy/refactor/improved-remote-claude-code
refactor: use new SessionService and SandboxService for remote Claude operations
2 parents b749fb9 + c957f1e commit 06671eb

File tree

13 files changed

+3751
-150
lines changed

13 files changed

+3751
-150
lines changed

‎pkg/api/rpc.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,11 +48,27 @@ func NewAgentClient() agentv1connect.AgentServiceClient {
4848
return agentv1connect.NewAgentServiceClient(getHTTPClient(getBaseURL()), getBaseURL(), WithUserAgent())
4949
}
5050

51+
func NewSessionClient() agentv1connect.SessionServiceClient {
52+
return agentv1connect.NewSessionServiceClient(getHTTPClient(getBaseURL()), getBaseURL(), WithUserAgent())
53+
}
54+
55+
func NewSandboxClient() agentv1connect.SandboxServiceClient {
56+
return agentv1connect.NewSandboxServiceClient(getHTTPClient(getBaseURL()), getBaseURL(), WithUserAgent())
57+
}
58+
5159
func WithAuthentication[T any](req *connect.Request[T], token string) *connect.Request[T] {
5260
req.Header().Add("Authorization", "Bearer "+token)
5361
return req
5462
}
5563

64+
func WithAuthenticationAndOrg[T any](req *connect.Request[T], token, orgID string) *connect.Request[T] {
65+
req.Header().Add("Authorization", "Bearer "+token)
66+
if orgID != "" {
67+
req.Header().Add("x-depot-org", orgID)
68+
}
69+
return req
70+
}
71+
5672
// getHTTPClient returns an HTTP client configured for the given base URL.
5773
// If the URL uses HTTP (not HTTPS), it configures h2c support for HTTP/2 over cleartext.
5874
func getHTTPClient(baseURL string) *http.Client {

‎pkg/cmd/claude/claude.go‎

Lines changed: 54 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ func NewCmdClaude() *cobra.Command {
3232
orgID string
3333
token string
3434
resumeSessionID string
35-
sandboxID string
3635
output string
3736
local bool
3837
repository *string
@@ -105,6 +104,18 @@ Subcommands:
105104
}
106105
ctx := cmd.Context()
107106

107+
// Check if we're running inside a sandbox
108+
sandboxID := os.Getenv("DEPOT_SANDBOX_ID")
109+
if sandboxID != "" {
110+
// Force local mode when running inside a sandbox
111+
local = true
112+
defer func() {
113+
if err := shutdownSandbox(ctx, sandboxID); err != nil {
114+
fmt.Fprintf(os.Stderr, "Warning: failed to shutdown sandbox: %v\n", err)
115+
}
116+
}()
117+
}
118+
108119
claudeArgs := []string{}
109120
for i := 0; i < len(args); i++ {
110121
arg := args[i]
@@ -156,20 +167,11 @@ Subcommands:
156167
} else {
157168
return fmt.Errorf("--resume flag requires a session ID")
158169
}
159-
case "--sandbox-id":
160-
if i+1 < len(args) && !strings.HasPrefix(args[i+1], "-") {
161-
sandboxID = args[i+1]
162-
i++
163-
} else {
164-
return fmt.Errorf("--sandbox-id flag requires a sandbox ID")
165-
}
166170
default:
167171
if strings.HasPrefix(arg, "--session-id=") {
168172
sessionID = strings.TrimPrefix(arg, "--session-id=")
169173
} else if strings.HasPrefix(arg, "--resume=") {
170174
resumeSessionID = strings.TrimPrefix(arg, "--resume=")
171-
} else if strings.HasPrefix(arg, "--sandbox-id=") {
172-
sandboxID = strings.TrimPrefix(arg, "--sandbox-id=")
173175
} else if strings.HasPrefix(arg, "--org=") {
174176
orgID = strings.TrimPrefix(arg, "--org=")
175177
} else if strings.HasPrefix(arg, "--token=") {
@@ -233,6 +235,11 @@ Subcommands:
233235
}
234236

235237
if !local {
238+
// Check if we're already in a sandbox - prevent recursive remote sessions
239+
if sandboxID != "" {
240+
return fmt.Errorf("cannot start a remote session from within a sandbox")
241+
}
242+
236243
// Default repository to current git remote only if --repository flag was specified with empty value
237244
var repoValue string
238245
if repository != nil {
@@ -255,12 +262,10 @@ Subcommands:
255262
Branch: branch,
256263
GitSecret: gitSecret,
257264
ResumeSessionID: resumeSessionID,
258-
RemoteSessionID: sandboxID,
259265
Wait: wait,
260266
Stdin: os.Stdin,
261267
Stdout: os.Stdout,
262268
Stderr: os.Stderr,
263-
AgentType: "claude",
264269
}
265270
return RunAgentRemote(ctx, agentOpts)
266271
} else {
@@ -284,24 +289,20 @@ Subcommands:
284289
}
285290

286291
// returns the session file UUID that claude should resume from
287-
func resumeSession(ctx context.Context, client agentv1connect.ClaudeServiceClient, token, sessionID, sessionDir, cwd, orgID string, retryCount int, retryDelay time.Duration) (string, error) {
288-
var resp *connect.Response[agentv1.DownloadClaudeSessionResponse]
292+
func resumeSession(ctx context.Context, client agentv1connect.SessionServiceClient, token, sessionID, sessionDir, cwd, orgID string, retryCount int, retryDelay time.Duration) (string, error) {
293+
var resp *connect.Response[agentv1.DownloadSessionResponse]
289294
var lastErr error
290295

291296
for i := range retryCount {
292297
if i > 0 {
293298
time.Sleep(retryDelay)
294299
}
295300

296-
req := &agentv1.DownloadClaudeSessionRequest{
297-
SessionId: sessionID,
298-
OrganizationId: new(string),
299-
}
300-
if orgID != "" {
301-
req.OrganizationId = &orgID
301+
req := &agentv1.DownloadSessionRequest{
302+
SessionId: sessionID,
302303
}
303304

304-
resp, lastErr = client.DownloadClaudeSession(ctx, api.WithAuthentication(connect.NewRequest(req), token))
305+
resp, lastErr = client.DownloadSession(ctx, api.WithAuthenticationAndOrg(connect.NewRequest(req), token, orgID))
305306
if lastErr == nil {
306307
break
307308
}
@@ -319,7 +320,8 @@ func resumeSession(ctx context.Context, client agentv1connect.ClaudeServiceClien
319320

320321
reader := bytes.NewReader(resp.Msg.SessionData)
321322

322-
sessionFilePath := filepath.Join(projectDir, fmt.Sprintf("%s.jsonl", resp.Msg.ClaudeSessionId))
323+
claudeSessionID := resp.Msg.ToolSessionId
324+
sessionFilePath := filepath.Join(projectDir, fmt.Sprintf("%s.jsonl", claudeSessionID))
323325
out, err := os.Create(sessionFilePath)
324326
if err != nil {
325327
return "", fmt.Errorf("failed to create session file: %w", err)
@@ -330,10 +332,10 @@ func resumeSession(ctx context.Context, client agentv1connect.ClaudeServiceClien
330332
return "", fmt.Errorf("failed to write session file: %w", err)
331333
}
332334

333-
return resp.Msg.ClaudeSessionId, nil
335+
return claudeSessionID, nil
334336
}
335337

336-
func saveSession(ctx context.Context, client agentv1connect.ClaudeServiceClient, token, sessionID, sessionFilePath string, retryCount int, retryDelay time.Duration, orgID string) error {
338+
func saveSession(ctx context.Context, client agentv1connect.SessionServiceClient, token, sessionID, sessionFilePath string, retryCount int, retryDelay time.Duration, orgID string) error {
337339
data, err := os.ReadFile(sessionFilePath)
338340
if err != nil {
339341
return fmt.Errorf("failed to read session file: %w", err)
@@ -349,21 +351,18 @@ func saveSession(ctx context.Context, client agentv1connect.ClaudeServiceClient,
349351

350352
claudeSessionID := filepath.Base(strings.TrimSuffix(sessionFilePath, ".jsonl"))
351353

352-
req := &agentv1.UploadClaudeSessionRequest{
353-
SessionData: data,
354-
SessionId: sessionID,
355-
OrganizationId: new(string),
356-
Summary: new(string),
357-
ClaudeSessionId: claudeSessionID,
354+
req := &agentv1.UploadSessionRequest{
355+
SessionData: data,
356+
SessionId: sessionID,
357+
Summary: new(string),
358+
ToolSessionId: claudeSessionID,
359+
AgentType: agentv1.AgentType_AGENT_TYPE_CLAUDE_CODE,
358360
}
359361
if summary != "" {
360362
req.Summary = &summary
361363
}
362-
if orgID != "" {
363-
req.OrganizationId = &orgID
364-
}
365364

366-
_, err := client.UploadClaudeSession(ctx, api.WithAuthentication(connect.NewRequest(req), token))
365+
_, err := client.UploadSession(ctx, api.WithAuthenticationAndOrg(connect.NewRequest(req), token, orgID))
367366
if err != nil {
368367
lastErr = err
369368
continue
@@ -422,7 +421,7 @@ func convertPathToProjectName(path string) string {
422421
}
423422

424423
// continuouslySaveSessionFile monitors the project directory for new or changed session files and automatically saves them
425-
func continuouslySaveSessionFile(ctx context.Context, projectDir string, client agentv1connect.ClaudeServiceClient, token, sessionID, orgID string) error {
424+
func continuouslySaveSessionFile(ctx context.Context, projectDir string, client agentv1connect.SessionServiceClient, token, sessionID, orgID string) error {
426425
if err := os.MkdirAll(projectDir, 0755); err != nil {
427426
return fmt.Errorf("failed to create project directory: %w", err)
428427
}
@@ -526,7 +525,7 @@ func RunClaudeSession(ctx context.Context, opts *ClaudeSessionOptions) error {
526525
opts.OrgID = os.Getenv("DEPOT_ORG_ID")
527526
}
528527

529-
client := api.NewClaudeClient()
528+
client := api.NewSessionClient()
530529

531530
// early auth check to prevent starting Claude if saving or resuming will fail
532531
if err := verifyAuthentication(ctx, client, token, opts.OrgID); err != nil {
@@ -647,16 +646,29 @@ func extractSummaryFromSession(data []byte) string {
647646

648647
// verifyAuthentication performs an early auth check by calling the list-sessions API
649648
// this prevents starting Claude if authentication or organization access will fail
650-
func verifyAuthentication(ctx context.Context, client agentv1connect.ClaudeServiceClient, token, orgID string) error {
651-
req := &agentv1.ListClaudeSessionsRequest{}
652-
if orgID != "" {
653-
req.OrganizationId = &orgID
654-
}
649+
func verifyAuthentication(ctx context.Context, client agentv1connect.SessionServiceClient, token, orgID string) error {
650+
req := &agentv1.ListSessionsRequest{}
655651

656-
_, err := client.ListClaudeSessions(ctx, api.WithAuthentication(connect.NewRequest(req), token))
652+
_, err := client.ListSessions(ctx, api.WithAuthenticationAndOrg(connect.NewRequest(req), token, orgID))
657653
if err != nil {
658654
return fmt.Errorf("authentication failed: %w", err)
659655
}
660656

661657
return nil
662658
}
659+
660+
// shutdownSandbox calls the Shutdown API to gracefully terminate and snapshot the sandbox
661+
func shutdownSandbox(ctx context.Context, sandboxID string) error {
662+
client := api.NewSandboxClient()
663+
664+
req := &agentv1.ShutdownRequest{
665+
SandboxId: sandboxID,
666+
}
667+
668+
_, err := client.Shutdown(ctx, connect.NewRequest(req))
669+
if err != nil {
670+
return fmt.Errorf("failed to shutdown sandbox: %w", err)
671+
}
672+
673+
return nil
674+
}

‎pkg/cmd/claude/listsessions.go‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -61,11 +61,11 @@ In interactive mode, pressing Enter on a session will start Claude with that ses
6161
orgID = os.Getenv("DEPOT_ORG_ID")
6262
}
6363

64-
client := api.NewClaudeClient()
64+
client := api.NewSessionClient()
6565

6666
isInteractive := output == "" && helpers.IsTerminal()
6767

68-
var allSessions []*agentv1.ClaudeSession
68+
var allSessions []*agentv1.Session
6969
currentPageToken := pageToken
7070
maxPages := 5
7171
pagesLoaded := 0
@@ -74,15 +74,15 @@ In interactive mode, pressing Enter on a session will start Claude with that ses
7474
reqCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
7575
defer cancel()
7676

77-
req := &agentv1.ListClaudeSessionsRequest{}
78-
if orgID != "" {
79-
req.OrganizationId = &orgID
77+
agentType := agentv1.AgentType_AGENT_TYPE_CLAUDE_CODE
78+
req := &agentv1.ListSessionsRequest{
79+
AgentType: &agentType,
8080
}
8181
if currentPageToken != "" {
8282
req.PageToken = &currentPageToken
8383
}
8484

85-
resp, err := client.ListClaudeSessions(reqCtx, api.WithAuthentication(connect.NewRequest(req), token))
85+
resp, err := client.ListSessions(reqCtx, api.WithAuthenticationAndOrg(connect.NewRequest(req), token, orgID))
8686
if err != nil {
8787
return fmt.Errorf("failed to list sessions: %w", err)
8888
}
@@ -166,7 +166,7 @@ In interactive mode, pressing Enter on a session will start Claude with that ses
166166
return cmd
167167
}
168168

169-
func chooseSession(sessions []*agentv1.ClaudeSession) (*agentv1.ClaudeSession, error) {
169+
func chooseSession(sessions []*agentv1.Session) (*agentv1.Session, error) {
170170
if len(sessions) == 0 {
171171
fmt.Println("No Claude sessions found")
172172
return nil, nil
@@ -211,7 +211,7 @@ func chooseSession(sessions []*agentv1.ClaudeSession) (*agentv1.ClaudeSession, e
211211
}
212212

213213
type sessionItem struct {
214-
session *agentv1.ClaudeSession
214+
session *agentv1.Session
215215
summary string
216216
}
217217

@@ -246,8 +246,8 @@ func (i sessionItem) FilterValue() string {
246246

247247
type sessionListModel struct {
248248
list list.Model
249-
sessions []*agentv1.ClaudeSession
250-
choice *agentv1.ClaudeSession
249+
sessions []*agentv1.Session
250+
choice *agentv1.Session
251251
ctrlC bool
252252
}
253253

@@ -283,7 +283,7 @@ func (m sessionListModel) View() string {
283283
return docStyle.Render(m.list.View())
284284
}
285285

286-
func sessionWriteJSON(sessions []*agentv1.ClaudeSession) error {
286+
func sessionWriteJSON(sessions []*agentv1.Session) error {
287287
type sessionJSON struct {
288288
SessionID string `json:"session_id"`
289289
Summary string `json:"summary,omitempty"`
@@ -313,7 +313,7 @@ func sessionWriteJSON(sessions []*agentv1.ClaudeSession) error {
313313
return encoder.Encode(jsonSessions)
314314
}
315315

316-
func sessionWriteCSV(sessions []*agentv1.ClaudeSession) error {
316+
func sessionWriteCSV(sessions []*agentv1.Session) error {
317317
w := csv.NewWriter(os.Stdout)
318318
if err := w.Write([]string{"SESSION_ID", "SUMMARY", "UPDATED_AT", "CREATED_AT"}); err != nil {
319319
return err

0 commit comments

Comments
 (0)