Repository: GoogleCloudPlatform/kubectl-ai
File: pkg/agent/conversation.go (Go)

Code before:
		sessionID := parts[1]

		manager, err := sessions.NewSessionManager()
		if err != nil {
			return "", false, fmt.Errorf("failed to create session manager: %w", err)
		}

		session, err := manager.FindSessionByID(sessionID)
		if err != nil {
			return "", false, fmt.Errorf("failed to get session: %w", err)
		}

		c.sessionMu.Lock()
		c.ChatMessageStore = session
		c.session.ChatMessageStore = session
		c.session.Messages = session.ChatMessages()
		metadata, err := session.LoadMetadata()
		if err != nil {
			c.sessionMu.Unlock()
			return "", false, fmt.Errorf("failed to load session metadata: %w", err)
		}
		c.session.ID = session.ID
		c.session.CreatedAt = metadata.CreatedAt
		c.session.LastModified = metadata.LastAccessed
		c.llmChat.Initialize(c.session.ChatMessageStore.ChatMessages())
		c.sessionMu.Unlock()

		return fmt.Sprintf("Resumed session %s.", sessionID), true, nil
	}

	return "", false, nil
}

func (c *Agent) listModels(ctx context.Context) ([]string, error) {
	if c.availableModels == nil {
		modelNames, err := c.LLM.ListModels(ctx)
		if err != nil {
			return nil, fmt.Errorf("listing models: %w", err)
		}
		c.availableModels = modelNames
	}
	return c.availableModels, nil
}

func (c *Agent) DispatchToolCalls(ctx context.Context) error {
	log := klog.FromContext(ctx)
	// execute all pending function calls
	for _, call := range c.pendingFunctionCalls {
		// Only show "Running" message and proceed with execution for non-interactive commands
		toolDescription := call.ParsedToolCall.Description()


Review comment (refactor):
Functionality looks great. I think we can abstract out `loadSession(...)` functionality in some common place so that it can be called here and also in `agent.Init()`. what do you think ?

Code after:
			return "", false, err
		}
		return fmt.Sprintf("Resumed session %s.", sessionID), true, nil
	}

	return "", false, nil
}

// loadSession loads a session by ID (or latest), updates the agent's state, and re-initializes the chat.
func (c *Agent) loadSession(sessionID string) error {
	manager, err := sessions.NewSessionManager()
	if err != nil {
		return fmt.Errorf("failed to create session manager: %w", err)
	}

	var session *sessions.Session
	if sessionID == "" || sessionID == "latest" {
		s, err := manager.GetLatestSession()
		if err != nil {
			return fmt.Errorf("failed to get latest session: %w", err)
		}
		if s == nil {
			// This can happen if GetLatestSession returns nil, nil (no sessions exist)
			return fmt.Errorf("no sessions found to resume")
		}
		session = s
	} else {
		s, err := manager.FindSessionByID(sessionID)
		if err != nil {
			return fmt.Errorf("failed to get session %q: %w", sessionID, err)
		}
		session = s
	}

	c.sessionMu.Lock()
	defer c.sessionMu.Unlock()

	c.ChatMessageStore = session
	c.session.ChatMessageStore = session
	c.session.Messages = session.ChatMessages()
	metadata, err := session.LoadMetadata()
	if err != nil {
		return fmt.Errorf("failed to load session metadata: %w", err)
	}
	c.session.ID = session.ID
	c.session.CreatedAt = metadata.CreatedAt
	c.session.LastModified = metadata.LastAccessed

	if c.llmChat != nil {
		if err := c.llmChat.Initialize(c.session.ChatMessageStore.ChatMessages()); err != nil {
			return fmt.Errorf("failed to re-initialize chat with new session: %w", err)